def __init__(
self,
decoder_length: int,
layer_sizes: List[int],
static_outputs_per_time_step: int,
) -> None:
super().__init__()
self.decoder_length = decoder_length
self.static_outputs_per_time_step = static_outputs_per_time_step
with self.name_scope():
self.mlp = MLP(layer_sizes, flatten=False)
self.expander = nn.Dense(units=decoder_length *
static_outputs_per_time_step)
def __init__(self, layer_sizes: List[int], **kwargs) -> None:
super().__init__(**kwargs)
self.model = MLP(layer_sizes, flatten=True)