def call(self, inputs, params=(), **kwargs):
rngs = _pop_rng_and_split(kwargs, len(inputs))
result = [
self._layer(x, params=params, rng=r, **kwargs)
for x, r in zip(inputs, rngs)
]
return tuple(result)
def forward(self, inputs, params=(), state=(), **kwargs):
rngs = _pop_rng_and_split(kwargs, len(inputs))
results = [
self._layer(x, params=params, state=state, rng=r, **kwargs)
for x, r in zip(inputs, rngs)
]
# TODO(kitaev): think about how to merge state across copies in the map.
return tuple(results), self._layer.state
def call(self, inputs, params=(), state=(), **kwargs):
rngs = _pop_rng_and_split(kwargs, len(inputs))
results = [
self._layer(x, params=params, state=state, rng=r, **kwargs)
for x, r in zip(inputs, rngs)
]
result_outputs, result_states = zip(*results)
return tuple(result_outputs), tuple(result_states)