ott.neural.networks.velocity_field.VelocityField.create_train_state# VelocityField.create_train_state(rng, optimizer, input_dim, condition_dim=None)[source]# Create the training state. Parameters: rng (Array) – Random number generator. optimizer (Union[Array, ndarray, bool_, number, Iterable[ArrayTree], Mapping[Any, ArrayTree]]) – Optimizer. input_dim (int) – Dimensionality of the velocity field. condition_dim (Optional[int]) – Dimensionality of the condition of the velocity field. Return type: TrainState Returns: The training state.
