(self, *, carry, forward_state)
| 44 | return jnp.zeros([batch_size], dtype=self.dtype()) |
| 45 | |
| 46 | def forward(self, *, carry, forward_state): |
| 47 | logging.info("_Layer: carry=%s forward_state=%s", shapes(carry), shapes(forward_state)) |
| 48 | forward_state = forward_state + carry |
| 49 | self.add_summary("carry_mean", jnp.mean(carry)) |
| 50 | self.add_module_output("state", forward_state) |
| 51 | self.add_state_update("inc", 2) |
| 52 | return carry + self.parameters["inc"], forward_state |
| 53 | |
| 54 | |
| 55 | class _ComplicatedLayer(BaseLayer): |
nothing calls this directly
no test coverage detected