ott.neural.models.MLP.create_train_state