(carry_i, x_i)
| 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 |
nothing calls this directly
no test coverage detected