zyx 0.17.0

// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only WITH Classpath-exception-2.0

//! `DType` and constant

#![allow(unused)]

use crate::scalar::{bf16, f8e4m3, f8e5m2, f16};
use crate::{
    Scalar, ZyxError,
    kernel::{BOp, IDX_T, UOp},
    shape::Dim,
};
use nanoserde::{DeBin, SerBin};
use std::fmt::{Debug, Display};

/// Represents the data type used for operations.
#[cfg_attr(feature = "py", pyo3::pyclass(eq, eq_int, from_py_object))]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, SerBin, DeBin)]
pub enum DType {
    /// 16 bit bfloat data type.
    BF16,
    /// 16 bit float data type.
    F16,
    /// 32 bit float data type.
    F32,
    /// 64 bit float data type.
    F64,
    /// 8 bit float data type, 4 exponent bits, 3 mantissa bits.
    F8E4M3,
    /// 8 bit float data type, 5 exponent bits, 2 mantissa bits.
    F8E5M2,
    /// 8 bit unsigned integer data type.
    U8,
    /// 16 bit unsigned integer data type.
    U16,
    /// 32 bit unsigned integer data type.
    U32,
    /// 64 bit unsigned integer data type.
    U64,
    /// 8 bit signed integer data type.
    I8,
    /// 16 bit signed integer data type.
    I16,
    /// 32 bit signed integer data type.
    I32,
    /// 64 bit signed integer data type.
    I64,
    /// 8 bit boolean data type.
    Bool,
}

impl DType {
    pub(crate) const N_DTYPES: usize = 15;
}

#[derive(Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, SerBin, DeBin)]
pub enum Constant {
    BF16([u8; 2]), // le bytes
    F16([u8; 2]),  // le bytes
    F32([u8; 4]),  // le bytes
    F64([u8; 8]),  // le bytes
    F8E4M3(u8),
    F8E5M2(u8),
    U8(u8),
    U16(u16),
    U32(u32),
    U64([u8; 8]), // le bytes
    I8(i8),
    I16(i16),
    I32(i32),
    I64([u8; 8]), // le bytes
    Bool(bool),
}

impl Display for DType {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.write_str(match self {
            Self::BF16 => "bf16",
            Self::F16 => "f16",
            Self::F32 => "f32",
            Self::F64 => "f64",
            Self::F8E4M3 => "f8e4m3",
            Self::F8E5M2 => "f8e5m2",
            Self::U8 => "u8",
            Self::U16 => "u16",
            Self::U32 => "u32",
            Self::U64 => "u64",
            Self::I8 => "i8",
            Self::I16 => "i16",
            Self::I32 => "i32",
            Self::I64 => "i64",
            Self::Bool => "bool",
        })
    }
}

impl DType {
    /// Is this dtype floating point?
    #[must_use]
    pub const fn is_float(self) -> bool {
        use DType::*;
        match self {
            BF16 | F16 | F32 | F64 | F8E4M3 | F8E5M2 => true,
            U8 | U16 | U32 | U64 | I8 | I16 | I32 | I64 | Bool => false,
        }
    }

    /// Is this dtype integer?
    #[must_use]
    pub const fn is_int(self) -> bool {
        use DType::*;
        match self {
            BF16 | F16 | F32 | F64 | F8E4M3 | F8E5M2 | Bool => false,
            U8 | U16 | U32 | U64 | I8 | I16 | I32 | I64 => true,
        }
    }

    /// Is this dtype unsigned integer?
    #[must_use]
    pub const fn is_uint(self) -> bool {
        use DType::*;
        match self {
            BF16 | F16 | F32 | F64 | Bool | I8 | I16 | I32 | I64 => false,
            F8E4M3 | F8E5M2 => false,
            U8 | U16 | U32 | U64 => true,
        }
    }

    /// Is `self -> dtype` cast value-preserving for every input?
    ///
    /// Used by Cast-of-Cast folding: collapsing `Cast(Cast(a→B)→C)` into
    /// `Cast(a→C)` is only valid when `A→B` is lossless. When in doubt a
    /// pair is absent here (lossy): that only misses an optimization,
    /// never miscompiles. Note `F16<->BF16` is lossy in both directions
    /// (mantissa vs range), and float->int is never lossless.
    #[must_use]
    pub fn is_lossless_cast(self, dtype: DType) -> bool {
        use DType::*;
        if self == dtype {
            return true;
        }
        match self {
            F8E4M3 | F8E5M2 => matches!(dtype, F16 | BF16 | F32 | F64),
            F16 => matches!(dtype, F32 | F64),
            BF16 => matches!(dtype, F32 | F64),
            F32 => matches!(dtype, F64),
            F64 => false,
            U8 => matches!(dtype, U16 | U32 | U64 | I16 | I32 | I64 | F16 | BF16 | F32 | F64),
            U16 => matches!(dtype, U32 | U64 | I32 | I64 | F32 | F64),
            U32 => matches!(dtype, U64 | I64 | F64),
            U64 => false,
            I8 => matches!(dtype, I16 | I32 | I64 | U16 | U32 | U64 | F16 | BF16 | F32 | F64),
            I16 => matches!(dtype, I32 | I64 | U32 | U64 | F32 | F64),
            I32 => matches!(dtype, I64 | U64 | F64),
            I64 => false,
            Bool => matches!(dtype, U8 | U16 | U32 | U64 | I8 | I16 | I32 | I64 | F8E4M3 | F8E5M2 | F16 | BF16 | F32 | F64),
        }
    }

    pub(crate) fn least_upper_dtype(self, rhs: DType) -> DType {
        use DType::*;
        // define an ordered list of "widening" priority
        let order = [
            Bool, U8, U16, U32, U64, I8, I16, I32, I64, F8E4M3, F8E5M2, BF16, F16, F32, F64,
        ];

        let i1 = order.iter().position(|&d| d == self).unwrap();
        let i2 = order.iter().position(|&d| d == rhs).unwrap();
        order[i1.max(i2)]
    }

    /// Get the size of this dtype in bits
    #[must_use]
    pub const fn bit_size(&self) -> u8 {
        use DType::*;
        match self {
            U8 | I8 | Bool | F8E4M3 | F8E5M2 => 8,
            BF16 | F16 | I16 | U16 => 16,
            F32 | I32 | U32 => 32,
            F64 | I64 | U64 => 64,
        }
    }

    /// Get initial constant for initializing accumulator for reduction op
    #[must_use]
    pub(super) fn init_for_rop(self, rop: BOp) -> Constant {
        match rop {
            BOp::Add => self.zero_constant(),
            BOp::Sub => unreachable!("Sub is not associative, cannot be a reduction"),
            BOp::Mul => self.one_constant(),
            BOp::Div => unreachable!("Div is not associative, cannot be a reduction"),
            BOp::Pow => unreachable!("Pow is not associative, cannot be a reduction"),
            BOp::Mod => unreachable!("Mod is not associative, cannot be a reduction"),
            BOp::Cmplt => unreachable!("comparison, cannot be a reduction"),
            BOp::Cmpgt => unreachable!("comparison, cannot be a reduction"),
            BOp::Cmpge => unreachable!("comparison, cannot be a reduction"),
            BOp::Max => self.min_constant(),
            BOp::Or => self.zero_constant(),
            BOp::And => self.one_constant(),
            BOp::BitXor => self.zero_constant(),
            BOp::BitOr => self.zero_constant(),
            BOp::BitAnd => self.zero_constant().unary(UOp::BitNot),
            BOp::BitShiftLeft => unreachable!("shift, cannot be a reduction"),
            BOp::BitShiftRight => unreachable!("shift, cannot be a reduction"),
            BOp::NotEq => unreachable!("comparison, cannot be a reduction"),
            BOp::Eq => unreachable!("comparison, cannot be a reduction"),
        }
    }

    #[must_use]
    pub(crate) const fn zero_constant(self) -> Constant {
        match self {
            Self::BF16 => Constant::BF16(bf16::ZERO.to_le_bytes()),
            Self::F16 => Constant::F16(f16::ZERO.to_le_bytes()),
            Self::F32 => Constant::F32(0f32.to_le_bytes()),
            Self::F64 => Constant::F64(0f64.to_le_bytes()),
            Self::F8E4M3 => Constant::F8E4M3(f8e4m3::ZERO.to_bits()),
            Self::F8E5M2 => Constant::F8E5M2(f8e5m2::ZERO.to_bits()),
            Self::U8 => Constant::U8(0),
            Self::U16 => Constant::U16(0),
            Self::U32 => Constant::U32(0),
            Self::I8 => Constant::I8(0),
            Self::I16 => Constant::I16(0),
            Self::I32 => Constant::I32(0),
            Self::I64 => Constant::I64(0i64.to_le_bytes()),
            Self::U64 => Constant::U64(0i64.to_le_bytes()),
            Self::Bool => Constant::Bool(false),
        }
    }

    #[must_use]
    pub(super) const fn one_constant(self) -> Constant {
        match self {
            Self::BF16 => Constant::BF16(bf16::ONE.to_le_bytes()),
            Self::F16 => Constant::F16(f16::ONE.to_le_bytes()),
            Self::F32 => Constant::F32(1f32.to_le_bytes()),
            Self::F64 => Constant::F64(1f64.to_le_bytes()),
            Self::F8E4M3 => Constant::F8E4M3(f8e4m3::ONE.to_bits()),
            Self::F8E5M2 => Constant::F8E5M2(f8e5m2::ONE.to_bits()),
            Self::U8 => Constant::U8(1),
            Self::U16 => Constant::U16(1),
            Self::U32 => Constant::U32(1),
            Self::I8 => Constant::I8(1),
            Self::I16 => Constant::I16(1),
            Self::I32 => Constant::I32(1),
            Self::I64 => Constant::I64(1i64.to_le_bytes()),
            Self::U64 => Constant::U64(1i64.to_le_bytes()),
            Self::Bool => Constant::Bool(true),
        }
    }

    #[must_use]
    pub(super) const fn min_constant(self) -> Constant {
        match self {
            Self::BF16 => Constant::BF16(bf16::MIN.to_le_bytes()),
            Self::F16 => Constant::F16(f16::MIN.to_le_bytes()),
            Self::F32 => Constant::F32(f32::MIN.to_le_bytes()),
            Self::F64 => Constant::F64(f64::MIN.to_le_bytes()),
            Self::F8E4M3 => Constant::F8E4M3(f8e4m3::MIN.to_bits()),
            Self::F8E5M2 => Constant::F8E5M2(f8e5m2::MIN.to_bits()),
            Self::U8 => Constant::U8(u8::MIN),
            Self::U16 => Constant::U16(u16::MIN),
            Self::U32 => Constant::U32(u32::MIN),
            Self::I8 => Constant::I8(i8::MIN),
            Self::I16 => Constant::I16(i16::MIN),
            Self::I32 => Constant::I32(i32::MIN),
            Self::I64 => Constant::I64(i64::MIN.to_le_bytes()),
            Self::U64 => Constant::U64(u64::MIN.to_le_bytes()),
            Self::Bool => Constant::Bool(false),
        }
    }

    #[must_use]
    pub(super) const fn safetensors(&self) -> &str {
        match self {
            Self::BF16 => "BF16",
            Self::F16 => "F16",
            Self::F32 => "F32",
            Self::F64 => "F64",
            Self::F8E4M3 => "F8_E4M3",
            Self::F8E5M2 => "F8_E5M2",
            Self::U8 => "U8",
            Self::U16 => "U16",
            Self::U32 => "U32",
            Self::I8 => "I8",
            Self::I16 => "I16",
            Self::I32 => "I32",
            Self::I64 => "I64",
            Self::U64 => "U64",
            Self::Bool => "BOOL",
        }
    }

    pub(super) fn from_safetensors(text: &str) -> Result<Self, ZyxError> {
        Ok(match text {
            "BF16" => Self::BF16,
            "F16" => Self::F16,
            "F32" => Self::F32,
            "F64" => Self::F64,
            "F8_E4M3" => Self::F8E4M3,
            "F8_E5M2" => Self::F8E5M2,
            "U8" => Self::U8,
            "U16" => Self::U16,
            "U32" => Self::U32,
            "I8" => Self::I8,
            "I16" => Self::I16,
            "I32" => Self::I32,
            "I64" => Self::I64,
            "U64" => Self::U64,
            "BOOL" => Self::Bool,
            _ => {
                return Err(ZyxError::ParseError(format!("Could not parse dtype {text}").into()));
            }
        })
    }
}

impl Constant {
    pub(crate) fn new<T: Scalar>(x: T) -> Self {
        use core::mem::transmute_copy as t;
        match T::dtype() {
            DType::BF16 => Self::BF16(unsafe { t(&x) }),
            DType::F16 => Self::F16(unsafe { t(&x) }),
            DType::F32 => Self::F32(unsafe { t(&x) }),
            DType::F64 => Self::F64(unsafe { t(&x) }),
            DType::F8E4M3 => Self::F8E4M3(unsafe { t(&x) }),
            DType::F8E5M2 => Self::F8E5M2(unsafe { t(&x) }),
            DType::U8 => Self::U8(unsafe { t(&x) }),
            DType::U16 => Self::U16(unsafe { t(&x) }),
            DType::U32 => Self::U32(unsafe { t(&x) }),
            DType::U64 => Self::U64(unsafe { t(&x) }),
            DType::I8 => Self::I8(unsafe { t(&x) }),
            DType::I16 => Self::I16(unsafe { t(&x) }),
            DType::I32 => Self::I32(unsafe { t(&x) }),
            DType::I64 => Self::I64(unsafe { t(&x) }),
            DType::Bool => Self::Bool(unsafe { t(&x) }),
        }
    }

    #[must_use]
    pub(crate) fn to_le_bytes(self) -> Vec<u8> {
        match self {
            Constant::BF16(x) | Constant::F16(x) => x.to_vec(),
            Constant::F32(x) => x.to_vec(),
            Constant::F64(x) => x.to_vec(),
            Constant::U8(x) => vec![x],
            Constant::F8E4M3(x) | Constant::F8E5M2(x) => vec![x],
            Constant::U16(x) => x.to_le_bytes().to_vec(),
            Constant::U32(x) => x.to_le_bytes().to_vec(),
            Constant::U64(x) | Constant::I64(x) => x.to_vec(),
            Constant::I8(x) => vec![x as u8],
            Constant::I16(x) => x.to_le_bytes().to_vec(),
            Constant::I32(x) => x.to_le_bytes().to_vec(),
            Constant::Bool(x) => vec![x as u8],
        }
    }

    #[allow(unused)]
    #[must_use]
    pub(crate) fn as_dim(self) -> Option<Dim> {
        match self {
            Constant::U8(d) => Some(Dim::from(d)),
            Constant::U16(d) => Some(Dim::from(d)),
            Constant::U32(d) => Some(Dim::from(d)),
            Constant::U64(d) => Some(i64::from_le_bytes(d)),
            Constant::I8(d) => {
                if d >= 0 {
                    Some(d as Dim)
                } else {
                    None
                }
            }
            Constant::I16(d) => {
                if d >= 0 {
                    Some(d as Dim)
                } else {
                    None
                }
            }
            Constant::I32(d) => {
                if d >= 0 {
                    Some(d as Dim)
                } else {
                    None
                }
            }
            Constant::I64(d) => {
                let d = i64::from_le_bytes(d);
                if d >= 0 { Some(d as Dim) } else { None }
            }
            Constant::Bool(d) => Some(Dim::from(d)),
            _ => None,
        }
    }

    #[must_use]
    pub(crate) fn is_max(self) -> bool {
        match self {
            Constant::U8(d) => d == u8::MAX,
            Constant::U16(d) => d == u16::MAX,
            Constant::U32(d) => d == u32::MAX,
            Constant::U64(d) => u64::from_le_bytes(d) == u64::MAX,
            Constant::I8(d) => d == i8::MAX,
            Constant::I16(d) => d == i16::MAX,
            Constant::I32(d) => d == i32::MAX,
            Constant::I64(d) => i64::from_le_bytes(d) == i64::MAX,
            _ => false,
        }
    }

    // Create new constant for indexing purposes
    pub(crate) fn idx<T: Scalar>(idx: T) -> Self {
        match IDX_T {
            DType::U32 => Self::U32(idx.cast::<u32>()),
            DType::U64 => Self::U64(idx.cast::<u64>().to_le_bytes()),
            DType::I64 => Self::I64(idx.cast::<i64>().to_le_bytes()),
            x => unreachable!("{x}"),
        }
    }

    pub(crate) fn from_le_bytes(bytes: &[u8], dtype: DType) -> Self {
        match dtype {
            DType::BF16 => Self::BF16([bytes[0], bytes[1]]),
            DType::F16 => Self::F16([bytes[0], bytes[1]]),
            DType::F32 => Self::F32([bytes[0], bytes[1], bytes[2], bytes[3]]),
            DType::F64 => Self::F64([bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7]]),
            DType::U8 => Self::U8(u8::from_le_bytes([bytes[0]])),
            DType::F8E4M3 => Self::F8E4M3(bytes[0]),
            DType::F8E5M2 => Self::F8E5M2(bytes[0]),
            DType::U16 => Self::U16(u16::from_le_bytes([bytes[0], bytes[1]])),
            DType::U32 => Self::U32(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])),
            DType::U64 => Self::U64([bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7]]),
            DType::I8 => Self::I8(i8::from_le_bytes([bytes[0]])),
            DType::I16 => Self::I16(i16::from_le_bytes([bytes[0], bytes[1]])),
            DType::I32 => Self::I32(i32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])),
            DType::I64 => Self::I64([bytes[0], bytes[1], bytes[2], bytes[3], bytes[4], bytes[5], bytes[6], bytes[7]]),
            DType::Bool => Self::Bool(bytes[0] != 0),
        }
    }

    pub(crate) const fn dtype(&self) -> DType {
        match self {
            Self::BF16(_) => DType::BF16,
            Self::F16(_) => DType::F16,
            Self::F32(_) => DType::F32,
            Self::F64(_) => DType::F64,
            Self::F8E4M3(_) => DType::F8E4M3,
            Self::F8E5M2(_) => DType::F8E5M2,
            Self::U8(_) => DType::U8,
            Self::U16(_) => DType::U16,
            Self::U32(_) => DType::U32,
            Self::U64(_) => DType::U64,
            Self::I8(_) => DType::I8,
            Self::I16(_) => DType::I16,
            Self::I32(_) => DType::I32,
            Self::I64(_) => DType::I64,
            Self::Bool(_) => DType::Bool,
        }
    }

    pub(crate) fn is_positive(&self) -> bool {
        match *self {
            Constant::BF16(x) => bf16::from_le_bytes(x) >= bf16::ZERO,
            Constant::F16(x) => f16::from_le_bytes(x) >= f16::ZERO,
            Constant::F32(x) => f32::from_le_bytes(x) >= 0f32,
            Constant::F64(x) => f64::from_le_bytes(x) >= 0f64,
            Constant::F8E4M3(x) => f8e4m3::from_bits(x) >= f8e4m3::ZERO,
            Constant::F8E5M2(x) => f8e5m2::from_bits(x) >= f8e5m2::ZERO,
            Constant::U8(_) | Constant::U16(_) | Constant::U32(_) | Constant::U64(_) | Constant::Bool(_) => true,
            Constant::I8(x) => x >= 0,
            Constant::I16(x) => x >= 0,
            Constant::I32(x) => x >= 0,
            Constant::I64(x) => i64::from_le_bytes(x) >= 0,
        }
    }

    pub(crate) fn is_minimum(&self) -> bool {
        #[allow(clippy::float_cmp)]
        match *self {
            Constant::BF16(x) => bf16::from_le_bytes(x) == bf16::MIN,
            Constant::F16(x) => f16::from_le_bytes(x) == f16::MIN,
            Constant::F32(x) => f32::from_le_bytes(x) == f32::MIN,
            Constant::F64(x) => f64::from_le_bytes(x) == f64::MIN,
            Constant::F8E4M3(x) => f8e4m3::from_bits(x) == f8e4m3::MIN,
            Constant::F8E5M2(x) => f8e5m2::from_bits(x) == f8e5m2::MIN,
            Constant::U8(x) => x == u8::MIN,
            Constant::U16(x) => x == u16::MIN,
            Constant::U32(x) => x == u32::MIN,
            Constant::U64(x) => u64::from_le_bytes(x) == u64::MIN,
            Constant::I8(x) => x == i8::MIN,
            Constant::I16(x) => x == i16::MIN,
            Constant::I32(x) => x == i32::MIN,
            Constant::I64(x) => i64::from_le_bytes(x) == i64::MIN,
            Constant::Bool(x) => !x,
        }
    }

    pub(crate) fn is_zero(&self) -> bool {
        match *self {
            Constant::BF16(x) => bf16::from_le_bytes(x) == bf16::ZERO,
            Constant::F16(x) => f16::from_le_bytes(x) == f16::ZERO,
            Constant::F32(x) => f32::from_le_bytes(x) == 0f32,
            Constant::F64(x) => f64::from_le_bytes(x) == 0f64,
            Constant::F8E4M3(x) => f8e4m3::from_bits(x) == f8e4m3::ZERO,
            Constant::F8E5M2(x) => f8e5m2::from_bits(x) == f8e5m2::ZERO,
            Constant::U8(x) => x == 0,
            Constant::U16(x) => x == 0,
            Constant::U32(x) => x == 0,
            Constant::U64(x) => u64::from_le_bytes(x) == 0,
            Constant::I8(x) => x == 0,
            Constant::I16(x) => x == 0,
            Constant::I32(x) => x == 0,
            Constant::I64(x) => i64::from_le_bytes(x) == 0,
            Constant::Bool(x) => !x,
        }
    }

    #[allow(clippy::float_cmp)]
    pub(crate) fn is_one(&self) -> bool {
        match *self {
            Constant::BF16(x) => bf16::from_le_bytes(x) == bf16::ONE,
            Constant::F16(x) => f16::from_le_bytes(x) == f16::ONE,
            Constant::F32(x) => f32::from_le_bytes(x) == 1f32,
            Constant::F64(x) => f64::from_le_bytes(x) == 1f64,
            Constant::F8E4M3(x) => f8e4m3::from_bits(x) == f8e4m3::ONE,
            Constant::F8E5M2(x) => f8e5m2::from_bits(x) == f8e5m2::ONE,
            Constant::U8(x) => x == 1,
            Constant::U16(x) => x == 1,
            Constant::U32(x) => x == 1,
            Constant::U64(x) => u64::from_le_bytes(x) == 1,
            Constant::I8(x) => x == 1,
            Constant::I16(x) => x == 1,
            Constant::I32(x) => x == 1,
            Constant::I64(x) => i64::from_le_bytes(x) == 1,
            Constant::Bool(x) => x,
        }
    }

    #[allow(clippy::float_cmp)]
    pub(crate) fn is_two(&self) -> bool {
        match *self {
            Constant::BF16(x) => bf16::from_le_bytes(x) == bf16::ONE + bf16::ONE,
            Constant::F16(x) => f16::from_le_bytes(x) == f16::ONE + f16::ONE,
            Constant::F32(x) => f32::from_le_bytes(x) == 2f32,
            Constant::F64(x) => f64::from_le_bytes(x) == 2f64,
            Constant::F8E4M3(x) => f8e4m3::from_bits(x) == f8e4m3::ONE + f8e4m3::ONE,
            Constant::F8E5M2(x) => f8e5m2::from_bits(x) == f8e5m2::ONE + f8e5m2::ONE,
            Constant::U8(x) => x == 2,
            Constant::U16(x) => x == 2,
            Constant::U32(x) => x == 2,
            Constant::U64(x) => u64::from_le_bytes(x) == 2,
            Constant::I8(x) => x == 2,
            Constant::I16(x) => x == 2,
            Constant::I32(x) => x == 2,
            Constant::I64(x) => i64::from_le_bytes(x) == 2,
            Constant::Bool(_) => false,
        }
    }

    pub(crate) const fn is_power_of_two(&self) -> bool {
        match *self {
            Constant::U32(x) => x != 0 && x.is_power_of_two(),
            Constant::U64(x) => {
                let x = u64::from_le_bytes(x);
                x != 0 && x.is_power_of_two()
            }
            _ => false,
        }
    }

    /// Bit-reinterpret `self` as `dtype`: raw bits are copied, never
    /// value-converted (unlike [`Constant::cast`]). Requires equal bit widths
    /// and `debug_assert`s it.
    #[must_use]
    pub(crate) fn bitcast(self, dtype: DType) -> Constant {
        debug_assert_eq!(self.dtype().bit_size(), dtype.bit_size(), "bitcast requires equal bit widths");
        let mut bytes = [0u8; 8];
        match self {
            Constant::BF16(x) | Constant::F16(x) => bytes[..2].copy_from_slice(&x),
            Constant::F32(x) => bytes[..4].copy_from_slice(&x),
            Constant::F64(x) | Constant::U64(x) | Constant::I64(x) => bytes.copy_from_slice(&x),
            Constant::U8(x) => bytes[0] = x,
            Constant::F8E4M3(x) | Constant::F8E5M2(x) => bytes[0] = x,
            Constant::I8(x) => bytes[0] = x as u8,
            Constant::U16(x) => bytes[..2].copy_from_slice(&x.to_le_bytes()),
            Constant::I16(x) => bytes[..2].copy_from_slice(&x.to_le_bytes()),
            Constant::U32(x) => bytes[..4].copy_from_slice(&x.to_le_bytes()),
            Constant::I32(x) => bytes[..4].copy_from_slice(&x.to_le_bytes()),
            Constant::Bool(x) => bytes[0] = u8::from(x),
        }
        match dtype {
            DType::BF16 => Constant::BF16([bytes[0], bytes[1]]),
            DType::F16 => Constant::F16([bytes[0], bytes[1]]),
            DType::F32 => Constant::F32([bytes[0], bytes[1], bytes[2], bytes[3]]),
            DType::F64 => Constant::F64(bytes),
            DType::U8 => Constant::U8(bytes[0]),
            DType::F8E4M3 => Constant::F8E4M3(bytes[0]),
            DType::F8E5M2 => Constant::F8E5M2(bytes[0]),
            DType::U16 => Constant::U16(u16::from_le_bytes([bytes[0], bytes[1]])),
            DType::U32 => Constant::U32(u32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])),
            DType::U64 => Constant::U64(bytes),
            DType::I8 => Constant::I8(bytes[0] as i8),
            DType::I16 => Constant::I16(i16::from_le_bytes([bytes[0], bytes[1]])),
            DType::I32 => Constant::I32(i32::from_le_bytes([bytes[0], bytes[1], bytes[2], bytes[3]])),
            DType::I64 => Constant::I64(bytes),
            DType::Bool => Constant::Bool(bytes[0] != 0),
        }
    }

    pub(super) fn cast(self, dtype: DType) -> Constant {
        match self {
            Constant::BF16(x) => bf16::from_le_bytes(x).cast_dtype(dtype),
            Constant::F16(x) => f16::from_le_bytes(x).cast_dtype(dtype),
            Constant::F32(x) => f32::from_le_bytes(x).cast_dtype(dtype),
            Constant::F64(x) => f64::from_le_bytes(x).cast_dtype(dtype),
            Constant::F8E4M3(x) => f8e4m3::from_bits(x).cast_dtype(dtype),
            Constant::F8E5M2(x) => f8e5m2::from_bits(x).cast_dtype(dtype),
            Constant::U8(x) => x.cast_dtype(dtype),
            Constant::I8(x) => x.cast_dtype(dtype),
            Constant::I16(x) => x.cast_dtype(dtype),
            Constant::U16(x) => x.cast_dtype(dtype),
            Constant::U32(x) => x.cast_dtype(dtype),
            Constant::U64(x) => u64::from_le_bytes(x).cast_dtype(dtype),
            Constant::I32(x) => x.cast_dtype(dtype),
            Constant::I64(x) => i64::from_le_bytes(x).cast_dtype(dtype),
            Constant::Bool(x) => x.cast_dtype(dtype),
        }
    }

    pub(super) fn unary(self, uop: UOp) -> Constant {
        use crate::Float;
        if uop == UOp::Not {
            return match self {
                Constant::Bool(x) => Constant::Bool(!x),
                _ => unreachable!("Not is only supported for bool"),
            };
        }
        if uop == UOp::BitNot {
            return match self {
                Constant::U8(x) => Constant::U8(!x),
                Constant::U16(x) => Constant::U16(!x),
                Constant::U32(x) => Constant::U32(!x),
                Constant::U64(x) => Constant::U64((!u64::from_le_bytes(x)).to_le_bytes()),
                Constant::I8(x) => Constant::I8(!x),
                Constant::I16(x) => Constant::I16(!x),
                Constant::I32(x) => Constant::I32(!x),
                Constant::I64(x) => Constant::I64((!i64::from_le_bytes(x)).to_le_bytes()),
                Constant::Bool(x) => Constant::Bool(!x),
                _ => unreachable!("BitNot is not supported for float types"),
            };
        }
        fn unary_func<T: Scalar>(x: T, uop: UOp) -> T {
            match uop {
                UOp::Reciprocal
                | UOp::Sqrt
                | UOp::Rsqrt
                | UOp::Sin
                | UOp::Cos
                | UOp::Floor
                | UOp::Trunc
                | UOp::Abs
                | UOp::Exp => {
                    unreachable!()
                }
                UOp::BitNot => unreachable!(),
                UOp::Not => unreachable!(),
                UOp::Neg => x.neg(),
                UOp::Exp2 => x.exp2(),
                UOp::Log2 => x.log2(),
            }
        }
        fn unary_func_float<T: Float>(x: T, uop: UOp) -> T {
            match uop {
                UOp::BitNot => unreachable!(),
                UOp::Not => unreachable!(),
                UOp::Neg => x.neg(),
                UOp::Reciprocal => x.reciprocal(),
                UOp::Sqrt => x.sqrt(),
                UOp::Rsqrt => x.sqrt().reciprocal(),
                UOp::Sin => x.sin(),
                UOp::Cos => x.cos(),
                UOp::Floor => x.floor(),
                UOp::Trunc => x.trunc(),
                UOp::Abs => x.abs(),
                UOp::Exp => x.exp(),
                UOp::Exp2 => x.exp2(),
                UOp::Log2 => x.log2(),
            }
        }
        match self {
            Constant::BF16(x) => Constant::BF16(unary_func_float(bf16::from_le_bytes(x), uop).to_le_bytes()),
            Constant::F16(x) => Constant::F16(unary_func_float(f16::from_le_bytes(x), uop).to_le_bytes()),
            Constant::F32(x) => Constant::F32(unary_func_float(f32::from_le_bytes(x), uop).to_le_bytes()),
            Constant::F64(x) => Constant::F64(unary_func_float(f64::from_le_bytes(x), uop).to_le_bytes()),
            Constant::F8E4M3(x) => Constant::F8E4M3(unary_func_float(f8e4m3::from_bits(x), uop).to_bits()),
            Constant::F8E5M2(x) => Constant::F8E5M2(unary_func_float(f8e5m2::from_bits(x), uop).to_bits()),
            Constant::U8(x) => Constant::U8(unary_func(x, uop)),
            Constant::U16(x) => Constant::U16(unary_func(x, uop)),
            Constant::U32(x) => Constant::U32(unary_func(x, uop)),
            Constant::U64(x) => Constant::U64(unary_func(u64::from_le_bytes(x), uop).to_le_bytes()),
            Constant::I8(x) => Constant::I8(unary_func(x, uop)),
            Constant::I16(x) => Constant::I16(unary_func(x, uop)),
            Constant::I32(x) => Constant::I32(unary_func(x, uop)),
            Constant::I64(x) => Constant::I64(unary_func(i64::from_le_bytes(x), uop).to_le_bytes()),
            Constant::Bool(x) => Constant::Bool(unary_func(x, uop)),
        }
    }
    // Assumes both constants are the same dtype
    pub(super) fn binary(x: Constant, y: Constant, bop: BOp) -> Constant {
        fn binary_func<T: Scalar>(x: T, y: T, bop: BOp) -> Constant {
            match bop {
                BOp::Add => Constant::new(x.add(y)),
                BOp::Sub => Constant::new(x.sub(y)),
                BOp::Mul => Constant::new(x.mul(y)),
                BOp::Div => Constant::new(x.div(y)),
                BOp::Pow => Constant::new(x.pow(y)),
                BOp::Mod => Constant::new(x.mod_(y)),
                BOp::Max => Constant::new(x.max(y)),
                BOp::Cmplt => Constant::new(x.cmplt(y)),
                BOp::Cmpgt => Constant::new(x.cmpgt(y)),
                BOp::Cmpge => Constant::new(x.cmpgt(y) || x.is_equal(y)),
                BOp::Or => Constant::new(x.or(y)),
                BOp::And => Constant::new(x.and(y)),
                BOp::NotEq => Constant::new(x.noteq(y)),
                BOp::Eq => Constant::new(x.is_equal(y)),
                BOp::BitXor => Constant::new(x.bitxor(y)),
                BOp::BitOr => Constant::new(x.bitor(y)),
                BOp::BitAnd => Constant::new(x.bitand(y)),
                BOp::BitShiftLeft => Constant::new(x.bitshiftleft(y)),
                BOp::BitShiftRight => Constant::new(x.bitshiftright(y)),
            }
        }
        debug_assert_eq!(x.dtype(), y.dtype());
        let res = match x {
            Constant::BF16(x) => {
                let Constant::BF16(y) = y else { unreachable!() };
                binary_func(bf16::from_le_bytes(x), bf16::from_le_bytes(y), bop)
            }
            Constant::F16(x) => {
                let Constant::F16(y) = y else { unreachable!() };
                binary_func(f16::from_le_bytes(x), f16::from_le_bytes(y), bop)
            }
            Constant::F32(x) => {
                let Constant::F32(y) = y else { unreachable!() };
                binary_func(f32::from_le_bytes(x), f32::from_le_bytes(y), bop)
            }
            Constant::F64(x) => {
                let Constant::F64(y) = y else { unreachable!() };
                binary_func(f64::from_le_bytes(x), f64::from_le_bytes(y), bop)
            }
            Constant::F8E4M3(x) => {
                let Constant::F8E4M3(y) = y else { unreachable!() };
                binary_func(f8e4m3::from_bits(x), f8e4m3::from_bits(y), bop)
            }
            Constant::F8E5M2(x) => {
                let Constant::F8E5M2(y) = y else { unreachable!() };
                binary_func(f8e5m2::from_bits(x), f8e5m2::from_bits(y), bop)
            }
            Constant::U8(x) => {
                let Constant::U8(y) = y else { unreachable!() };
                binary_func(x, y, bop)
            }
            Constant::U16(x) => {
                let Constant::U16(y) = y else { unreachable!() };
                binary_func(x, y, bop)
            }
            Constant::U32(x) => {
                let Constant::U32(y) = y else { unreachable!() };
                binary_func(x, y, bop)
            }
            Constant::U64(x) => {
                let Constant::U64(y) = y else { unreachable!() };
                binary_func(u64::from_le_bytes(x), u64::from_le_bytes(y), bop)
            }
            Constant::I8(x) => {
                let Constant::I8(y) = y else { unreachable!() };
                binary_func(x, y, bop)
            }
            Constant::I16(x) => {
                let Constant::I16(y) = y else { unreachable!() };
                binary_func(x, y, bop)
            }
            Constant::I32(x) => {
                let Constant::I32(y) = y else { unreachable!() };
                binary_func(x, y, bop)
            }
            Constant::I64(x) => {
                let Constant::I64(y) = y else { unreachable!() };
                binary_func(i64::from_le_bytes(x), i64::from_le_bytes(y), bop)
            }
            Constant::Bool(x) => {
                let Constant::Bool(y) = y else { unreachable!() };
                binary_func(x, y, bop)
            }
        };
        // A folded NaN is always a bug: some rewrite proved a guard
        // degenerate and produced e.g. 0/0. Fail loudly at the fold site.
        if Self::is_nan_const(&res) {
            panic!("constant folding produced NaN: {bop:?}({x:?}, {y:?})");
        }
        res
    }

    fn is_nan_const(c: &Constant) -> bool {
        match c {
            Constant::F32(x) => f32::from_le_bytes(*x).is_nan(),
            Constant::F64(x) => f64::from_le_bytes(*x).is_nan(),
            Constant::BF16(x) => u16::from_le_bytes(*x) & 0x7fff > 0x7f80,
            Constant::F16(x) => {
                let b = u16::from_le_bytes(*x);
                b & 0x7c00 == 0x7c00 && b & 0x03ff != 0
            }
            Constant::F8E4M3(x) => *x == 0x7f || *x == 0xff,
            Constant::F8E5M2(x) => *x & 0x7c == 0x7c && *x & 0x03 != 0,
            _ => false,
        }
    }
}

trait CastDType: Scalar {
    fn cast_dtype(self, dtype: DType) -> Constant {
        match dtype {
            DType::BF16 => Constant::BF16(self.cast::<bf16>().to_le_bytes()),
            DType::F16 => Constant::F16(self.cast::<f16>().to_le_bytes()),
            DType::F32 => Constant::F32(self.cast::<f32>().to_le_bytes()),
            DType::F64 => Constant::F64(self.cast::<f64>().to_le_bytes()),
            DType::F8E4M3 => Constant::F8E4M3(self.cast::<f8e4m3>().to_bits()),
            DType::F8E5M2 => Constant::F8E5M2(self.cast::<f8e5m2>().to_bits()),
            DType::U8 => Constant::U8(self.cast()),
            DType::U16 => Constant::U16(self.cast()),
            DType::U32 => Constant::U32(self.cast()),
            DType::U64 => Constant::U64(self.cast::<u64>().to_le_bytes()),
            DType::I8 => Constant::I8(self.cast()),
            DType::I16 => Constant::I16(self.cast()),
            DType::I32 => Constant::I32(self.cast()),
            DType::I64 => Constant::I64(self.cast::<i64>().to_le_bytes()),
            DType::Bool => Constant::Bool(self.cast()),
        }
    }
}

impl<T: Scalar> CastDType for T {}

impl Display for Constant {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        match self {
            Self::BF16(value) => f.write_fmt(format_args!("{}", bf16::from_le_bytes(*value))),
            Self::F16(value) => f.write_fmt(format_args!("{}", f16::from_le_bytes(*value))),
            Self::F32(value) => f.write_fmt(format_args!("{}", f32::from_le_bytes(*value))),
            Self::F64(value) => f.write_fmt(format_args!("{}", f64::from_le_bytes(*value))),
            Self::F8E4M3(value) => f.write_fmt(format_args!("{}", f8e4m3::from_bits(*value))),
            Self::F8E5M2(value) => f.write_fmt(format_args!("{}", f8e5m2::from_bits(*value))),
            Self::U8(value) => f.write_fmt(format_args!("{value}")),
            Self::U16(value) => f.write_fmt(format_args!("{value}")),
            &Self::U64(value) => f.write_fmt(format_args!("{}", u64::from_le_bytes(value))),
            Self::U32(value) => f.write_fmt(format_args!("{value}")),
            Self::I8(value) => f.write_fmt(format_args!("{value}")),
            Self::I16(value) => f.write_fmt(format_args!("{value}")),
            Self::I32(value) => f.write_fmt(format_args!("{value}")),
            &Self::I64(value) => f.write_fmt(format_args!("{}", i64::from_le_bytes(value))),
            Self::Bool(value) => f.write_fmt(format_args!("{value}")),
        }
    }
}

impl Debug for Constant {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        f.write_fmt(format_args!("{self}"))
    }
}

/// Quantized data type of a GGUF super-block.
///
/// The numeric codes match the `ggml_type` enum (`ggml.h`): the same codes
/// [`Tensor::load_gguf`](crate::Tensor::load_gguf) matches on. Every variant
/// loads as raw `[num_blocks, block_bytes]` [`DType::U8`] via `load_gguf` and
/// decodes to dense floats via
/// [`Tensor::dequantize`](crate::Tensor::dequantize).
// Variant names mirror `ggml.h` (`Q4_0`, `Q4_K`, ...) instead of UpperCamelCase
// so they stay grep-compatible with the ggml docs and error messages.
#[allow(non_camel_case_types)]
#[cfg_attr(feature = "py", pyo3::pyclass(eq, eq_int))]
#[derive(Debug, Clone, Copy, PartialEq, Eq, PartialOrd, Ord, Hash, SerBin, DeBin)]
pub enum QDType {
    /// 4-bit, 32 elements in 18 bytes (`block_q4_0`).
    Q4_0,
    /// 4-bit, 32 elements in 20 bytes (`block_q4_1`).
    Q4_1,
    /// 5-bit, 32 elements in 22 bytes (`block_q5_0`).
    Q5_0,
    /// 5-bit, 32 elements in 24 bytes (`block_q5_1`).
    Q5_1,
    /// 8-bit, 32 elements in 34 bytes (`block_q8_0`).
    Q8_0,
    /// 8-bit, 32 elements in 36 bytes (`block_q8_1`).
    Q8_1,
    /// 2-bit super-block, 256 elements in 84 bytes (`block_q2_K`).
    Q2_K,
    /// 3-bit super-block, 256 elements in 110 bytes (`block_q3_K`).
    Q3_K,
    /// 4-bit super-block, 256 elements in 144 bytes (`block_q4_K`).
    Q4_K,
    /// 5-bit super-block, 256 elements in 176 bytes (`block_q5_K`).
    Q5_K,
    /// 6-bit super-block, 256 elements in 210 bytes (`block_q6_K`).
    Q6_K,
    /// 8-bit super-block, 256 elements in 292 bytes (`block_q8_K`).
    Q8_K,
    /// Improved 2-bit, 256 elements in 66 bytes (`block_iq2_xxs`).
    IQ2_XXS,
    /// Improved 2-bit, 256 elements in 74 bytes (`block_iq2_xs`).
    IQ2_XS,
    /// Improved 2-bit, 256 elements in 82 bytes (`block_iq2_s`).
    IQ2_S,
    /// Improved 3-bit, 256 elements in 98 bytes (`block_iq3_xxs`).
    IQ3_XXS,
    /// Improved 3-bit, 256 elements in 110 bytes (`block_iq3_s`).
    IQ3_S,
    /// Improved 1-bit, 256 elements in 50 bytes (`block_iq1_s`).
    IQ1_S,
    /// Improved 1-bit, 256 elements in 56 bytes (`block_iq1_m`).
    IQ1_M,
    /// Non-linear 4-bit, 32 elements in 18 bytes (`block_iq4_nl`).
    IQ4_NL,
    /// Non-linear 4-bit, 256 elements in 136 bytes (`block_iq4_xs`).
    IQ4_XS,
}

impl QDType {
    /// GGUF type code of this quant (`ggml.h` numbering).
    #[must_use]
    pub const fn gguf_code(self) -> u32 {
        match self {
            Self::Q4_0 => 2,
            Self::Q4_1 => 3,
            Self::Q5_0 => 6,
            Self::Q5_1 => 7,
            Self::Q8_0 => 8,
            Self::Q8_1 => 9,
            Self::Q2_K => 10,
            Self::Q3_K => 11,
            Self::Q4_K => 12,
            Self::Q5_K => 13,
            Self::Q6_K => 14,
            Self::Q8_K => 15,
            Self::IQ2_XXS => 16,
            Self::IQ2_XS => 17,
            Self::IQ3_XXS => 18,
            Self::IQ1_S => 19,
            Self::IQ4_NL => 20,
            Self::IQ3_S => 21,
            Self::IQ2_S => 22,
            Self::IQ4_XS => 23,
            Self::IQ1_M => 29,
        }
    }

    /// Number of elements one super-block decodes to.
    #[must_use]
    pub const fn elems_per_block(self) -> i64 {
        match self {
            Self::Q4_0 | Self::Q4_1 | Self::Q5_0 | Self::Q5_1 | Self::Q8_0 | Self::Q8_1 | Self::IQ4_NL => 32,
            Self::Q2_K
            | Self::Q3_K
            | Self::Q4_K
            | Self::Q5_K
            | Self::Q6_K
            | Self::Q8_K
            | Self::IQ2_XXS
            | Self::IQ2_XS
            | Self::IQ2_S
            | Self::IQ3_XXS
            | Self::IQ3_S
            | Self::IQ1_S
            | Self::IQ1_M
            | Self::IQ4_XS => 256,
        }
    }

    /// Size in bytes of one super-block on disk.
    #[must_use]
    pub const fn block_bytes(self) -> i64 {
        match self {
            Self::Q4_0 => 18,
            Self::Q4_1 => 20,
            Self::Q5_0 => 22,
            Self::Q5_1 => 24,
            Self::Q8_0 => 34,
            Self::Q8_1 => 36,
            Self::Q2_K => 84,
            Self::Q3_K => 110,
            Self::Q4_K => 144,
            Self::Q5_K => 176,
            Self::Q6_K => 210,
            Self::Q8_K => 292,
            Self::IQ2_XXS => 66,
            Self::IQ2_XS => 74,
            Self::IQ2_S => 82,
            Self::IQ3_XXS => 98,
            Self::IQ3_S => 110,
            Self::IQ1_S => 50,
            Self::IQ1_M => 56,
            Self::IQ4_NL => 18,
            Self::IQ4_XS => 136,
        }
    }

    /// Maps a GGUF type code to a [`QDType`], or `None` for dense codes
    /// (`F32`, `F16`, ints) and removed/unsupported codes.
    #[must_use]
    pub const fn from_code(code: u32) -> Option<Self> {
        match code {
            2 => Some(Self::Q4_0),
            3 => Some(Self::Q4_1),
            6 => Some(Self::Q5_0),
            7 => Some(Self::Q5_1),
            8 => Some(Self::Q8_0),
            9 => Some(Self::Q8_1),
            10 => Some(Self::Q2_K),
            11 => Some(Self::Q3_K),
            12 => Some(Self::Q4_K),
            13 => Some(Self::Q5_K),
            14 => Some(Self::Q6_K),
            15 => Some(Self::Q8_K),
            16 => Some(Self::IQ2_XXS),
            17 => Some(Self::IQ2_XS),
            18 => Some(Self::IQ3_XXS),
            19 => Some(Self::IQ1_S),
            20 => Some(Self::IQ4_NL),
            21 => Some(Self::IQ3_S),
            22 => Some(Self::IQ2_S),
            23 => Some(Self::IQ4_XS),
            29 => Some(Self::IQ1_M),
            _ => None,
        }
    }
}

impl Display for QDType {
    fn fmt(&self, f: &mut core::fmt::Formatter<'_>) -> core::fmt::Result {
        f.write_str(match self {
            Self::Q4_0 => "q4_0",
            Self::Q4_1 => "q4_1",
            Self::Q5_0 => "q5_0",
            Self::Q5_1 => "q5_1",
            Self::Q8_0 => "q8_0",
            Self::Q8_1 => "q8_1",
            Self::Q2_K => "q2_K",
            Self::Q3_K => "q3_K",
            Self::Q4_K => "q4_K",
            Self::Q5_K => "q5_K",
            Self::Q6_K => "q6_K",
            Self::Q8_K => "q8_K",
            Self::IQ2_XXS => "iq2_xxs",
            Self::IQ2_XS => "iq2_xs",
            Self::IQ2_S => "iq2_s",
            Self::IQ3_XXS => "iq3_xxs",
            Self::IQ3_S => "iq3_s",
            Self::IQ1_S => "iq1_s",
            Self::IQ1_M => "iq1_m",
            Self::IQ4_NL => "iq4_nl",
            Self::IQ4_XS => "iq4_xs",
        })
    }
}