cellflow.networks.ConditionalVelocityField.create_train_state

ConditionalVelocityField.create_train_state(rng, optimizer, input_dim, conditions)

Create the training state.

Parameters:

• rng (Array) – Random number generator.

• optimizer (Union[Array, ndarray, bool, number, bool, int, float, complex, Iterable[ArrayTree], Mapping[Any, ArrayTree]]) – Optimizer.

• input_dim (int) – Dimensionality of the velocity field.

• conditions (dict[str, Array]) – Conditions describing the perturbation.

Return type:

TrainState

Returns:

: The training state.