(params)
| 139 | ) |
| 140 | |
| 141 | def create_trainstate_from_params(params): |
| 142 | return TrainState.create(params=params, tx=optimizer, apply_fn=None) |
| 143 | |
| 144 | def init_fn(rng): |
| 145 | rng_generator = JaxRNG(rng) |
nothing calls this directly
no outgoing calls
no test coverage detected