def METHOD_NAME(self, batch_size, **kwargs):
    return {
        "additions": self.addition_factor(
            shape=(batch_size, 1, 1, 1), dtype=self.compute_dtype
        ),
        "thresholds": self.threshold_factor(
            shape=(batch_size, 1, 1, 1), dtype=self.compute_dtype
        ),
    }