Skip to main content

ExtensionType

burn::backend

Trait ExtensionType 

pub trait ExtensionType<B>
where
    B: Backend,
{
    type Target;

    // Required methods
    fn map_to_dispatch<F>(
        self,
        map_kind: F,
        autodiff: DispatchAutodiffContext,
    ) -> Self::Target
       where F: Fn(BackendTensor<B>) -> DispatchTensorKind;
    fn map_from_dispatch<F>(target: Self::Target, unwrap_kind: F) -> Self
       where F: Fn(DispatchTensor) -> BackendTensor<B>;
    fn routing_tensor(target: &Self::Target) -> Option<&DispatchTensor>;
    fn routing_float_tensor(target: &Self::Target) -> Option<&DispatchTensor>;
}
Expand description

A trait to map custom structs and enums of tensor primitives across the Dispatch boundary, in both directions.

This trait cooperates with the #[backend_extension] macro. When an extension operation returns such a type, map_to_dispatch wraps each internal tensor into a DispatchTensor; when it takes one as an input, map_from_dispatch reconstructs the concrete value and routing_tensor / routing_float_tensor locate a tensor for dispatch routing and backend selection. autodiff_context merges the contexts of all tensor fields participating in one operation. Nested #[extension_type] fields are traversed recursively.

Implementations are generated automatically using #[derive(ExtensionType)].

Required Associated Types§

type Target

The target struct layout where all internal concrete backend tensors are transformed into DispatchTensors.

Required Methods§

fn map_to_dispatch<F>( self, map_kind: F, autodiff: DispatchAutodiffContext, ) -> Self::Target
where F: Fn(BackendTensor<B>) -> DispatchTensorKind,

Transforms the internal fields of the struct by applying a backend-specific wrapping closure.

§Arguments
  • map_kind - A closure provided by the dispatch macro that knows how to map a backend-agnostic BackendTensor variant into the correct DispatchTensorKind variant (e.g., Wgpu, Cuda, Cpu).
  • autodiff - The semantic autodiff backend context to attach to each DispatchTensor.
§Returns

A new instance of the struct mapped to the Dispatch backend.

fn map_from_dispatch<F>(target: Self::Target, unwrap_kind: F) -> Self
where F: Fn(DispatchTensor) -> BackendTensor<B>,

Reconstruct the concrete Struct<B> from its dispatch form Struct<Dispatch>.

This is the inverse of map_to_dispatch, used when a custom struct is passed as an input to a backend extension operation. The dispatch glue has already selected the target backend B; unwrap_kind pulls the matching BackendTensor out of each field’s DispatchTensorKind, and the derived impl calls the right accessor (.float(), .int(), …) per field to recover the concrete primitive.

§Arguments
  • unwrap_kind - A closure provided by the dispatch macro that validates the tensor’s representation and unwraps its DispatchTensorKind into the BackendTensor for the selected backend B, panicking on a backend mismatch.

fn routing_tensor(target: &Self::Target) -> Option<&DispatchTensor>

Return a tensor of the dispatch form to use for routing, or None if this value currently holds no tensor (e.g. an enum on a tensor-less variant).

A struct/enum input carries no top-level DispatchTensor of its own, so the dispatch glue uses this to read the runtime backend tag (.kind) and autodiff context. All other tensor fields must carry the same context; dispatch validates them while mapping the value to its concrete backend. This lookup recurses into nested #[extension_type] fields.

fn routing_float_tensor(target: &Self::Target) -> Option<&DispatchTensor>

Like routing_tensor but returns only a float tensor, or None if there is none.

The dispatch glue prefers a float routing tensor because active float presence decides whether the operation needs an autodiff backend. The glue falls back to routing_tensor only when no float tensor exists anywhere in the inputs.

Dyn Compatibility§

This trait is not dyn compatible.

In older versions of Rust, dyn compatibility was called "object safety".

Implementors§