Skip to main content

burn_router/
lib.rs

1#![cfg_attr(not(feature = "std"), no_std)]
2#![warn(missing_docs)]
3#![cfg_attr(docsrs, feature(doc_cfg))]
4#![recursion_limit = "138"]
5
6//! Burn multi-backend router.
7
8mod backend;
9mod bridge;
10mod channel;
11mod client;
12mod ops;
13mod runner;
14mod tensor;
15mod types;
16
17pub use backend::*;
18pub use bridge::*;
19pub use channel::*;
20pub use client::*;
21pub use runner::*;
22pub use tensor::*;
23pub use types::*;
24
25/// A local channel with a simple byte bridge between backends.
26/// It transfers tensors between backends via the underlying [tensor data](burn_backend::TensorData).
27pub type DirectByteChannel<Backends> = DirectChannel<Backends, ByteBridge<Backends>>;
28
29/// Router backend.
30///
31/// # Example
32///
33/// ```ignore
34/// type MyBackend = Router<(NdArray, Wgpu)>;
35/// ```
36pub type Router<Backends> = BackendRouter<DirectByteChannel<Backends>>;
37
38extern crate alloc;
39
40#[cfg(test)]
41#[allow(unused)]
42mod tests {
43    use crate::BackendRouter;
44    use crate::DirectByteChannel;
45
46    pub type TestBackend1 = burn_ndarray::NdArray<f32, i32>;
47    pub type TestBackend2 = burn_wgpu::Wgpu<f32, i32>;
48    pub type TestBackend = BackendRouter<DirectByteChannel<(TestBackend1, TestBackend2)>>;
49}