To implement a custom RNN cell in ParaRNN, you must define three distinct components:
1. System Parameters
Create a dataclass inheriting from SystemParameters. This defines the learnable parameters (e.g., weights) and functions (e.g., activations) passed to the recurrence step. You must implement unpack() and repack() methods to allow the framework to manipulate parameters.
2. Implementation Class
Choose a base class based on your Jacobian structure to enable CUDA acceleration:
RNNCellDenseImpl: For dense Jacobians (no optimization).RNNCellDiagImpl: For diagonal Jacobians (enables parallel_CUDA).RNNCellBlockDiagImpl: For block-diagonal Jacobians (enables parallel_CUDA).
Implement the recurrence_step class method: recurrence_step(cls, x, h, system_parameters) where h_t = f(h_{t-1}, x_t; params).
3. Cell Class
Create a class inheriting from BaseRNNCell. This manages the state, handles parameter initialization in _specific_init, and provides the _system_parameters property.
Once defined, the cell automatically supports sequential and parallel modes. If you used a -Diag or -BlockDiag implementation, parallel_CUDA is also available out-of-the-box.
# 1. Define System Parameters
@dataclass
class MyRNNSystemParameters(SystemParameters[MyRNNTrait]):
my_weight: torch.Tensor
my_nonlin: typ.Callable
def unpack(self) -> typ.Tuple[torch.Tensor, typ.Callable]:
return (self.my_weight, self.my_nonlin)
@classmethod
def repack(cls: T, pars: typ.Tuple[torch.Tensor, typ.Callable]) -> T:
return MyRNNSystemParameters(my_weight=pars[0], my_nonlin=pars[1])
# 2. Choose Implementation
class MyRNNImpl(RNNCellDiagImpl[MyRNNSystemParameters]):
@classmethod
def recurrence_step(cls, x, h, system_parameters):
# Core recurrence logic here
...
# 3. Define Cell Class
class MyRNNCell(BaseRNNCell[MyRNNConfig, MyRNNSystemParameters, MyRNNImpl]):
def __init__(self, config):
super().__init__(config)
def _specific_init(self, config):
self.my_weight = torch.nn.Parameter(...)
self.my_nonlin = ...
@property
def _system_parameters(self):
return MyRNNSystemParameters(my_weight=self.my_weight, my_nonlin=self.my_nonlin)
# Usage
config = MyRNNConfig(state_dim=64, input_dim=32, mode='parallel')
model = MyRNNCell(config)
output = model(x)
