diff options
| author | Thomas Santerre <thomas@santerre.xyz> | 2024-04-04 16:32:47 -0400 |
|---|---|---|
| committer | GitHub <noreply@github.com> | 2024-04-04 22:32:47 +0200 |
| commit | c5626b827147e5029c6bd3e37352ec8ac501cfc3 (patch) | |
| tree | cabc1c5d8713c52159e503b87277410d1b776d9a /candle-metal-kernels | |
| parent | e6a5b82ba6507e7e21d5a5d45241bd8f005609b7 (diff) | |
| download | candle-c5626b827147e5029c6bd3e37352ec8ac501cfc3.tar.gz candle-c5626b827147e5029c6bd3e37352ec8ac501cfc3.tar.bz2 candle-c5626b827147e5029c6bd3e37352ec8ac501cfc3.zip | |
Add support for "sign" on tensors (#2012)
* add the sign unary operator
* remove uneeded import
* remove uneeded import
* undo formatting
* undo formatting
* remove unnecessary redefintion
* allow gradient to flow through for sign and round
* fix cpu ops to ensure that negzero and positive zero are handled properly
* clippy fixes
* Properly avoid gradient tracking.
* Use a branchless version.
---------
Co-authored-by: laurent <laurent.mazare@gmail.com>
Diffstat (limited to 'candle-metal-kernels')
| -rw-r--r-- | candle-metal-kernels/src/lib.rs | 2 | ||||
| -rw-r--r-- | candle-metal-kernels/src/unary.metal | 2 |
2 files changed, 3 insertions, 1 deletions
diff --git a/candle-metal-kernels/src/lib.rs b/candle-metal-kernels/src/lib.rs index 140927e3..5af48fae 100644 --- a/candle-metal-kernels/src/lib.rs +++ b/candle-metal-kernels/src/lib.rs @@ -193,7 +193,7 @@ macro_rules! ops{ pub mod unary { ops!( cos, sin, exp, sqr, sqrt, neg, log, gelu, abs, ceil, floor, relu, round, erf, gelu_erf, - tanh, recip, silu + tanh, recip, silu, sign ); } pub mod binary { diff --git a/candle-metal-kernels/src/unary.metal b/candle-metal-kernels/src/unary.metal index bdc13f9e..809522d7 100644 --- a/candle-metal-kernels/src/unary.metal +++ b/candle-metal-kernels/src/unary.metal @@ -145,6 +145,7 @@ UNARY_OP(erf) UNARY_OP(tanh) UNARY_OP(recip) UNARY_OP(relu) +UNARY_OP(sign) UNARY(id, float, copy_f32, copy_f32_strided) UNARY(id, half, copy_f16, copy_f16_strided) UNARY(id, uint8_t, copy_u8, copy_u8_strided) @@ -174,6 +175,7 @@ BFLOAT_UNARY_OP(erf) BFLOAT_UNARY_OP(tanh) BFLOAT_UNARY_OP(recip) BFLOAT_UNARY_OP(relu) +BFLOAT_UNARY_OP(sign) UNARY(id, bfloat, copy_bf16, copy_bf16_strided) |