burn_tensor/tensor/api/
transaction.rs

1use super::Tensor;
2use crate::{ExecutionError, TensorData};
3use alloc::vec::Vec;
4use burn_backend::ops::TransactionPrimitive;
5use burn_dispatch::Dispatch;
6
7/// A transaction can [read](Self::register) multiple tensors at once with a single operation improving
8/// compute utilization with optimized laziness.
9///
10/// # Example
11///
12/// ```rust,ignore
13///  let [output_data, loss_data, targets_data] = Transaction::default()
14///    .register(output)
15///    .register(loss)
16///    .register(targets)
17///    .execute()
18///    .try_into()
19///    .expect("Correct amount of tensor data");
20/// ```
21pub struct Transaction {
22    opaque: transaction_opaque::Opaque,
23}
24
25burn_std::obfuscate!(
26    type: TransactionPrimitive<Dispatch>,
27    module: transaction_opaque,
28    derives: [Send, Sync],
29);
30
31impl Default for Transaction {
32    fn default() -> Self {
33        Self::from_op(TransactionPrimitive::<Dispatch>::default())
34    }
35}
36
37impl Transaction {
38    /// Crate-internal constructor wrapping a dispatch-level transaction.
39    pub(crate) fn from_op(op: TransactionPrimitive<Dispatch>) -> Self {
40        Self {
41            opaque: transaction_opaque::Opaque::new(op),
42        }
43    }
44
45    /// Crate-internal mutable borrow of the underlying transaction primitive.
46    pub(crate) fn as_op_mut(&mut self) -> &mut TransactionPrimitive<Dispatch> {
47        self.opaque.as_mut()
48    }
49
50    /// Crate-internal owning extraction of the underlying transaction primitive.
51    pub(crate) fn into_op(self) -> TransactionPrimitive<Dispatch> {
52        self.opaque.into_inner()
53    }
54
55    /// Add a [tensor](Tensor) to the transaction to be read.
56    pub fn register<const D: usize, K: crate::kind::Transaction>(
57        mut self,
58        tensor: Tensor<D, K>,
59    ) -> Self {
60        K::register_transaction(self.as_op_mut(), tensor.primitive);
61        self
62    }
63
64    /// Executes the transaction synchronously and returns the [data](TensorData) in the same order
65    /// in which they were [registered](Self::register).
66    pub fn execute(self) -> Vec<TensorData> {
67        burn_std::future::block_on(self.execute_async())
68            .expect("Error while reading data: use `try_execute` to handle error at runtime")
69    }
70
71    /// Executes the transaction synchronously and returns the [data](TensorData) in the same
72    /// order in which they were [registered](Self::register).
73    ///
74    /// # Returns
75    ///
76    /// Any error that might have occurred since the last time the device was synchronized.
77    pub fn try_execute(self) -> Result<Vec<TensorData>, ExecutionError> {
78        burn_std::future::block_on(self.execute_async())
79    }
80
81    /// Executes the transaction asynchronously and returns the [data](TensorData) in the same order
82    /// in which they were [registered](Self::register).
83    pub async fn execute_async(self) -> Result<Vec<TensorData>, ExecutionError> {
84        self.into_op().execute_async().await
85    }
86}
burn_tensor/tensor/api/transaction.rs

burn_tensor/tensor/api/
transaction.rs
