Opaque container that is not traversed by jax.tree.map.
| 1866 | pps: Nested[_AdastarPerParamState] # pytype: disable=invalid-annotation |
| 1867 | |
| 1868 | class _AdastarUpdateResult(flax_struct.PyTreeNode): |
| 1869 | """Opaque container that is not traversed by jax.tree.map.""" |
| 1870 | |
| 1871 | updates: Tensor # the update to apply to params. |
| 1872 | pps: _AdastarPerParamState # pytype: disable=invalid-annotation |
| 1873 | |
| 1874 | update_schedule = schedule.as_schedule_fn(update_schedule) |
| 1875 |
no outgoing calls
no test coverage detected