burn_mamba/utils/
backend_macros.rs

1//! Macros to cut down the per-backend boilerplate for `*SsdBackendExt` traits.
2//!
3//! Both Mamba-2 and Mamba-3 define an SSD-backend extension trait whose default
4//! implementation already does the right thing for every burn backend except
5//! `Autodiff` (which gets the custom backward). The macros below emit the
6//! per-backend "use the default impl" blocks and the autodiff marker trait.
7
8/// Emit `impl $trait_name for <backend> {}` blocks for every burn backend
9/// supported by this crate, opting in to the trait's default body.
10///
11/// Each impl is feature-gated by the corresponding `backend-*` / `cubecl` /
12/// `fusion` feature.
13#[macro_export]
14macro_rules! impl_ssd_backend_ext_for_burn_backends {
15    ($trait_name:path) => {
16        #[cfg(feature = "backend-ndarray")]
17        impl<F, I> $trait_name for burn::backend::NdArray<F, I> {}
18
19        #[cfg(feature = "backend-flex")]
20        impl $trait_name for burn::backend::Flex {}
21
22        #[cfg(any(feature = "backend-tch-cpu", feature = "backend-tch-gpu"))]
23        impl<F, I> $trait_name for burn::backend::libtorch::LibTorch<F, I> {}
24
25        #[cfg(feature = "backend-remote")]
26        impl<F, I> $trait_name for burn::backend::RemoteBackend<F, I> {}
27
28        #[cfg(feature = "cubecl")]
29        impl<R> $trait_name for burn_cubecl::CubeBackend<R> where R: burn_cubecl::CubeRuntime {}
30
31        // Fusion delegates to the inner backend's default impl.
32        #[cfg(feature = "fusion")]
33        impl<B: burn_fusion::FusionBackend + $trait_name> $trait_name for burn_fusion::Fusion<B> {}
34    };
35}
36
37/// Declare a marker trait `$autodiff_trait: Backend + $ext_trait + AutodiffBackend`
38/// and a blanket impl for `burn::backend::Autodiff<B>` whenever `B: $ext_trait`.
39///
40/// Both Mamba-2 and Mamba-3 expose this marker to let callers bound their own
41/// generics on "autodiff-capable backend that also implements the SSD ext."
42#[macro_export]
43macro_rules! decl_ssd_autodiff_backend_ext {
44    ($autodiff_trait:ident, $ext_trait:path $(, $extra_bound:path)*) => {
45        /// Marker for an autodiff-capable backend that also implements the SSD
46        /// extension trait (so the custom memory-efficient backward is available).
47        #[cfg(feature = "autodiff")]
48        pub trait $autodiff_trait:
49            burn::backend::Backend
50            + burn::backend::AutodiffBackend
51            + $ext_trait
52            $(+ $extra_bound)*
53        {
54        }
55        /// Any autodiff-compatible backend whose inner backend implements the
56        /// ext trait satisfies the marker. The actual custom backward lives in
57        /// `super::backward` (a custom impl of the ext trait for `Autodiff<B>`).
58        #[cfg(feature = "autodiff")]
59        impl<B> $autodiff_trait
60            for burn::backend::Autodiff<B>
61            where
62                B: burn::backend::Backend + $ext_trait $(+ $extra_bound)*,
63                B: burn::backend::BackendTypes<FloatTensorPrimitive = burn::backend::DispatchTensor>,
64        {
65        }
66    };
67}
burn_mamba/utils/backend_macros.rs

burn_mamba/utils/
backend_macros.rs
