huggingface/candle
1
0
Fork 0
mirror of https://github.com/huggingface/candle.git synced 2025-06-19 03:54:56 +00:00
Code Issues Packages Projects Releases Wiki Activity

Use Map2 for binary ops.

Browse Source
This commit is contained in:
laurent
2023-06-29 10:09:15 +01:00
parent 83c7d660ca
commit c9c468e1aa
2 changed files with 32 additions and 76 deletions
Download Patch File Download Diff File Expand all files Collapse all files

14
candle-core/src/op.rs
View File

@ -54,11 +54,8 @@ pub(crate) trait UnaryOp {
pub(crate) trait BinaryOp {
const NAME: &'static str;
const KERNEL_BF16: &'static str;
const KERNEL_F16: &'static str;
const KERNEL_F32: &'static str;
const KERNEL_F64: &'static str;
const KERNEL_U32: &'static str;
const KERNEL: &'static str;
const V: Self;
fn bf16(v1: bf16, v2: bf16) -> bf16;
fn f16(v1: f16, v2: f16) -> f16;
fn f32(v1: f32, v2: f32) -> f32;
@ -85,11 +82,8 @@ macro_rules! bin_op {
($op:ident, $name: literal, $e: expr) => {
impl BinaryOp for $op {
const NAME: &'static str = $name;
const KERNEL_BF16: &'static str = concat!("b", $name, "_bf16");
const KERNEL_F16: &'static str = concat!("b", $name, "_f16");
const KERNEL_F32: &'static str = concat!("b", $name, "_f32");
const KERNEL_F64: &'static str = concat!("b", $name, "_f64");
const KERNEL_U32: &'static str = concat!("b", $name, "_u32");
const KERNEL: &'static str = concat!("b", $name);
const V: Self = $op;
fn bf16(v1: bf16, v2: bf16) -> bf16 {
$e(v1, v2)
}