How to implement interior mutability for tensor fields in a Module?

[qcynaut]

Hi everyone,

I'm new to the Burn framework and I'm trying to create a module where I need to mutate a tensor field during the forward pass. Here's what I'm trying to achieve:

#[derive(Debug, Module)]
pub struct MyModule<B: Backend> {
    tensors: Tensor<B, 1>, // <- I need this to be mutable
    // ... other fields
}

The Problem:

My forward pass function needs to modify the tensors field, but I'm having trouble implementing interior mutability. I tried using RefCell, but it doesn't work with the Module derive macro.

What I've Tried:

• Using RefCell<Tensor<B, 1>> - this doesn't compile with the Module derive

Questions:

1. Is there a recommended pattern in Burn for handling mutable state within a module?
2. Should I be approaching this differently? Perhaps modules in Burn are meant to be immutable and I should return updated state instead?
3. Are there any examples of modules that need to maintain and update internal state during forward passes?

Any guidance would be greatly appreciated! Thanks in advance for your help.

[laggui]

1. If you only need single-threaded execution, you could define forward with &mut self and mutate internal fields directly. (e.g. self.tensors = other_computed_tensor or self.tensors = self.tensors.slice_assign(slice, partial_update). But this won't work correctly in multi-threaded or multi-device training, where updates need to be thread-safe. For a general thread-safe approach, Burn has RunningState as used in BatchNorm.
   

   burn/crates/burn-core/src/module/param/running.rs

   Lines 48 to 58 in 9ff428a

   	/// A state that can be updated during the forward pass while being thread safe.
   	///
   	/// # Note
   	///
   	/// The state value is the average of all updates on all threads.
   	#[derive(Clone, Debug)]
   	pub struct RunningState<V> {
   	id: ParamId,
   	values: Arc<Mutex<HashMap<ThreadId, V>>>,
   	value: Arc<Mutex<V>>,
   	}

   
   Though you might want to adapt it given the average reduction strategy that is currently specirfic for batchnorm.
2. Modules in Burn can have mutable state, but you should be careful about thread safety (hence, the running state approach).
3. BatchNorm, as linked above 🙂

[laggui]

Whoops, just noticed this was closed. Well, I answered anyway 😅

[qcynaut]

Thank you so much for the detailed response! I really appreciate you taking the time to answer even after I closed the discussion.

I had found a temporary workaround earlier, which is why I closed it, but your solution with RunningState is exactly what I needed - I honestly hadn't come across it in my initial research of the framework. The BatchNorm example is super helpful as a reference.

I'll definitely look into implementing the thread-safe approach with RunningState, especially since I want to make sure my code works correctly in multi-threaded scenarios. This is much cleaner than what I was trying to do!

Thanks again for the guidance 🙂