(self, *, num_iters, xs, **kwargs)
| 853 | """Tests scan_in_context.""" |
| 854 | |
| 855 | def _invoke(self, *, num_iters, xs, **kwargs): |
| 856 | batch_size = 2 |
| 857 | |
| 858 | def fn(carry_i, x_i): |
| 859 | ctx = current_context() |
| 860 | assert ctx is not None |
| 861 | state_i = ctx.state |
| 862 | # Add a nested output for testing filtering. |
| 863 | ctx.add_module_output( |
| 864 | "nested", |
| 865 | dict( |
| 866 | with_carry=dict( |
| 867 | output=x_i + carry_i, |
| 868 | with_state=dict(output=x_i + carry_i + state_i), |
| 869 | ), |
| 870 | output=x_i, |
| 871 | ), |
| 872 | ) |
| 873 | ctx.add_summary("carry", WeightedSummary(carry_i.mean(), carry_i.size)) |
| 874 | return carry_i + 1, x_i + carry_i + state_i |
| 875 | |
| 876 | xs["xs"] = jnp.arange(num_iters, dtype=jnp.int32)[:, None] * jnp.ones( |
| 877 | batch_size, dtype=jnp.int32 |
| 878 | ) |
| 879 | return scan_in_context(fn, carry=jnp.zeros(1, dtype=jnp.int32), xs=xs, **kwargs) |
| 880 | |
| 881 | @contextlib.contextmanager |
| 882 | def _dummy_context(self): |
no test coverage detected