zyx 0.15.2

// Copyright (C) 2025 zk4x
// SPDX-License-Identifier: LGPL-3.0-only

//! `DType` and constant

#![allow(unused)]

use crate::{
    Scalar, ZyxError,
    kernel::{BOp, IDX_T, UOp},
    shape::Dim,
};
use half::{bf16, f16};
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 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,
}

#[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
    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::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 {
        match self {
            Self::BF16 | Self::F16 | Self::F32 | Self::F64 => true,
            Self::U8 | Self::U16 | Self::U32 | Self::U64 | Self::I8 | Self::I16 | Self::I32 | Self::I64 | Self::Bool => false,
        }
    }

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

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

    pub(crate) fn least_upper_dtype(self, rhs: DType) -> DType {
        use DType::{BF16, Bool, F16, F32, F64, I8, I16, I32, I64, U8, U16, U32, U64};
        // define an ordered list of "widening" priority
        let order = [Bool, U8, U16, U32, U64, I8, I16, I32, I64, 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 {
        match self {
            Self::U8 | Self::I8 | Self::Bool => 8,
            Self::BF16 | Self::F16 | Self::I16 | Self::U16 => 16,
            Self::F32 | Self::I32 | Self::U32 => 32,
            Self::F64 | Self::I64 | Self::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 => todo!(),
            BOp::Mul => self.one_constant(),
            BOp::Div => todo!(),
            BOp::Pow => todo!(),
            BOp::Mod => todo!(),
            BOp::Cmplt => todo!(),
            BOp::Cmpgt => todo!(),
            BOp::Max => self.min_constant(),
            BOp::Or => todo!(),
            BOp::And => todo!(),
            BOp::BitXor => todo!(),
            BOp::BitOr => todo!(),
            BOp::BitAnd => todo!(),
            BOp::BitShiftLeft => todo!(),
            BOp::BitShiftRight => todo!(),
            BOp::NotEq => todo!(),
            BOp::Eq => todo!(),
        }
    }

    #[must_use]
    pub(super) 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::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(0u64.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::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(1u64.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::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::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,
            "U8" => Self::U8,
            "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::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) }),
        }
    }

    #[allow(unused)]
    #[must_use]
    pub(super) 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(u64::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,
        }
    }

    // Create new constant for indexing purposes
    pub(crate) fn idx<T: Scalar>(idx: T) -> Self {
        use core::mem::transmute_copy as t;
        match T::dtype() {
            DType::I32 => {
                let idx: i32 = unsafe { t(&idx) };
                if IDX_T == DType::U64 {
                    Self::U64((idx as u64).to_le_bytes())
                } else {
                    Self::U32(idx as u32)
                }
            }
            DType::U32 => {
                let idx: u32 = unsafe { t(&idx) };
                if IDX_T == DType::U64 {
                    Self::U64(u64::from(idx).to_le_bytes())
                } else {
                    Self::U32(idx)
                }
            }
            DType::U64 => {
                let idx: u64 = unsafe { t(&idx) };
                if IDX_T == DType::U64 {
                    Self::U64(idx.to_le_bytes())
                } else {
                    Self::U32(idx as u32)
                }
            }
            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::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::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::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::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::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::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::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,
        }
    }

    pub(super) fn cast(self, dtype: DType) -> Constant {
        match self {
            Constant::BF16(x) => half::bf16::from_le_bytes(x).cast_dtype(dtype),
            Constant::F16(x) => half::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::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;
        fn unary_func<T: Scalar>(x: T, uop: UOp) -> T {
            match uop {
                UOp::Reciprocal | UOp::Sqrt | UOp::Sin | UOp::Cos | UOp::Floor | UOp::Trunc | UOp::Abs => {
                    unreachable!()
                }
                UOp::BitNot => todo!(),
                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 => todo!(),
                UOp::Neg => x.neg(),
                UOp::Reciprocal => x.reciprocal(),
                UOp::Sqrt => x.sqrt(),
                UOp::Sin => x.sin(),
                UOp::Cos => x.cos(),
                UOp::Floor => x.floor(),
                UOp::Trunc => x.trunc(),
                UOp::Abs => x.abs(),
                UOp::Exp2 => x.exp2(),
                UOp::Log2 => x.log2(),
            }
        }
        match self {
            Constant::BF16(x) => Constant::BF16(unary_func_float(half::bf16::from_le_bytes(x), uop).to_le_bytes()),
            Constant::F16(x) => Constant::F16(unary_func_float(half::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::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::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());
        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::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)
            }
        }
    }
}

trait CastDType: Scalar {
    fn cast_dtype(self, dtype: DType) -> Constant {
        match dtype {
            DType::BF16 => Constant::BF16(self.cast::<half::bf16>().to_le_bytes()),
            DType::F16 => Constant::F16(self.cast::<half::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::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::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}"))
    }
}