ott.neural.networks.potentials.PotentialMLP.create_train_state

PotentialMLP.create_train_state(rng, optimizer, input, **kwargs)

Create initial training state.

Parameters:

• rng (Array)

• optimizer (Union[Array, ndarray, bool_, number, Iterable[ArrayTree], Mapping[Any, ArrayTree]])

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

• kwargs (Any)

Return type:

PotentialTrainState