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

Custom ops with a single argument (#214)

Browse Source 
* Add the CustomOp1 trait.

* Add an example of custom op.

* Polish the custom op example.

* Add some backward pass test for custom ops.
This commit is contained in:
Laurent Mazare
2023-07-21 16:18:05 +02:00
committed by GitHub
parent b02229ce92
commit a6bcdfb269
8 changed files with 241 additions and 18 deletions
Download Patch File Download Diff File Expand all files Collapse all files

7
candle-core/src/lib.rs
View File

@ -33,13 +33,13 @@
//!
//! Rust is cool, and a lot of the HF ecosystem already has Rust crates [safetensors](https://github.com/huggingface/safetensors) and [tokenizers](https://github.com/huggingface/tokenizers)
mod backend;
pub mod backend;
mod backprop;
mod conv;
mod convert;
mod cpu_backend;
pub mod cpu_backend;
#[cfg(feature = "cuda")]
mod cuda_backend;
pub mod cuda_backend;
mod device;
pub mod display;
mod dtype;
@ -65,6 +65,7 @@ pub use dtype::{DType, WithDType};
pub use error::{Error, Result};
pub use indexer::IndexOp;
pub use layout::Layout;
pub use op::CustomOp1;
pub use shape::{Shape, D};
pub use storage::Storage;
pub use strided_index::{StridedBlocks, StridedIndex};