(params)
| 94 | kwargs = {k: _to_partitioned_transformation(v) for k, v in kwargs.items()} |
| 95 | |
| 96 | def init_fn(params): |
| 97 | return {k: v.init(params) for k, v in kwargs.items()} |
| 98 | |
| 99 | def update_fn( |
| 100 | updates: NestedTensor, state: dict[str, Any], params: NestedOptParam |
nothing calls this directly
no test coverage detected