ott.solvers.nn.icnn.ICNN.create_train_state

ICNN.create_train_state(rng, optimizer, input)

Create initial TrainState.

Return type

TrainState

Parameters

• rng (jax.Array) –

• optimizer (Union[jax.Array, Iterable[ArrayTree], Mapping[Any, ArrayTree]]) –

• input (Union[int, Tuple[int, ...]]) –