(self)
| 899 | self._invoke(num_iters=3, xs={}) |
| 900 | |
| 901 | def test_basic(self): |
| 902 | num_iters = 3 |
| 903 | |
| 904 | # Invoke with inherited state. In this case, the same state is used each iter. |
| 905 | with self._dummy_context(): |
| 906 | carry, ys = self._invoke(num_iters=num_iters, xs={}) |
| 907 | self.assertNestedEqual(jnp.array([num_iters], dtype=carry.dtype), carry) |
| 908 | self.assertNestedEqual( |
| 909 | jnp.array([[0, 10], [2, 12], [4, 14]], dtype=carry.dtype), |
| 910 | ys, |
| 911 | ) |
| 912 | |
| 913 | # Invoke with explicit state. In this case, the state is unrolled. |
| 914 | with self._dummy_context(): |
| 915 | carry, ys = self._invoke( |
| 916 | num_iters=num_iters, |
| 917 | xs={"state": jnp.ones([num_iters, 1], dtype=jnp.int32) * 10}, |
| 918 | ) |
| 919 | self.assertNestedEqual(jnp.array([num_iters], dtype=carry.dtype), carry) |
| 920 | self.assertNestedEqual( |
| 921 | jnp.array([[10, 10], [12, 12], [14, 14]], dtype=carry.dtype), |
| 922 | ys, |
| 923 | ) |
| 924 | |
| 925 | def test_drop_output(self): |
| 926 | num_iters = 3 |
nothing calls this directly
no test coverage detected