Introduction to Keras for Researchers

library(keras3)
library(tensorflow)

x <- tf$constant(rbind(c(5, 2), c(1, 3)))
print(x)

as.array(x)

x$dtype

x$shape

tf$ones(shape = shape(2, 1))

tf$zeros(shape = shape(2, 1))

x <- random_normal(shape = c(2, 2), mean = 0.0, stddev = 1.0)
x <- random_uniform(shape = c(2, 2), minval = 0, maxval = 10)

initial_value <- random_normal(shape=c(2, 2))
a <- tf$Variable(initial_value)
print(a)

new_value <- random_normal(shape=c(2, 2))
a$assign(new_value)

added_value <- random_normal(shape=c(2, 2))
a$assign_add(added_value)

a <- random_normal(shape=c(2, 2))
b <- random_normal(shape=c(2, 2))

c <- a + b
d <- tf$square(c)
e <- tf$exp(d)

a <- random_normal(shape=c(2, 2))
b <- random_normal(shape=c(2, 2))

with(tf$GradientTape() %as% tape, {
  tape$watch(a)  # Start recording the history of operations applied to `a`
  c <- tf$sqrt(tf$square(a) + tf$square(b))  # Do some math using `a`
  # What's the gradient of `c` with respect to `a`?
  dc_da <- tape$gradient(c, a)
  print(dc_da)
})

a <- tf$Variable(a)

with(tf$GradientTape() %as% tape, {
  c <- tf$sqrt(tf$square(a) + tf$square(b))
  dc_da <- tape$gradient(c, a)
  print(dc_da)
})

with(tf$GradientTape() %as% outer_tape, {
  with(tf$GradientTape() %as% tape, {
    c <- tf$sqrt(tf$square(a) + tf$square(b))
    dc_da <- tape$gradient(c, a)
  })
  d2c_da2 <- outer_tape$gradient(dc_da, a)
  print(d2c_da2)
})

Linear <- new_layer_class(
  "Linear",
  initialize = function(units = 32, input_dim = 32) {
    super$initialize()
    self$w <- self$add_weight(
      shape = shape(input_dim, units),
      initializer = "random_normal",
      trainable = TRUE
    )
    self$b <- self$add_weight(
      shape = shape(units),
      initializer = "zeros",
      trainable = TRUE
    )
  },
  call = function(inputs) {
    tf$matmul(inputs, self$w) + self$b
  }
)

# Instantiate our layer.
linear_layer <- Linear(units=4, input_dim=2)

# The layer can be treated as a function.
# Here we call it on some data.
y <- linear_layer(tf$ones(shape(2, 2)))

linear_layer$weights

Linear <- new_layer_class(
  "Linear",
  initialize = function(units = 32) {
    super$initialize()
    self$units <- units
  },
  build = function(input_shape) {
    self$w <- self$add_weight(
      shape = shape(input_shape[-1], self$units),
      initializer = "random_normal",
      trainable = TRUE
    )
    self$b <- self$add_weight(
      shape = shape(self$units),
      initializer = "zeros",
      trainable = TRUE
    )
  },
  call = function(inputs) {
    tf$matmul(inputs, self$w) + self$b
  }
)

# Instantiate our layer.
linear_layer <- Linear(units = 4)

# This will also call `build(input_shape)` and create the weights.
y <- linear_layer(tf$ones(shape(2, 2)))

# Prepare a dataset.
c(c(x_train, y_train), .) %<-% dataset_mnist()

x_train <- array_reshape(x_train, c(60000, 784)) / 255

dataset <- tfdatasets::tensor_slices_dataset(list(x_train, y_train)) %>%
  tfdatasets::dataset_shuffle(buffer_size=1024) %>%
  tfdatasets::dataset_batch(64)

# Instantiate our linear layer (defined above) with 10 units.
linear_layer <- Linear(units = 10)

# Instantiate a logistic loss function that expects integer targets.
loss_fn <- loss_sparse_categorical_crossentropy(from_logits=TRUE)

# Instantiate an optimizer.
optimizer <- optimizer_sgd(learning_rate=1e-3)

# Iterate over the batches of the dataset.
coro::loop(for(data in dataset) {
  # Open a GradientTape.
  with(tf$GradientTape() %as% tape, {
    # Forward pass.
    logits <- linear_layer(data[[1]])

    # Loss value for this batch.
    loss_value <- loss_fn(data[[2]], logits)
  })

  # Get gradients of the loss wrt the weights.
  gradients <- tape$gradient(loss_value, linear_layer$trainable_weights)

  # Update the weights of our linear layer.
  optimizer$apply_gradients(zip_lists(gradients, linear_layer$trainable_weights))
})
loss_value

ComputeSum <- new_layer_class(
  "ComputeSum",
  initialize = function(input_dim) {
    super$initialize()
    # Create a non-trainable weight.
    self$total <- self$add_weight(
      initializer = "zeros",
      shape = shape(input_dim),
      trainable = FALSE
    )
  },
  call = function(inputs) {
    self$total$assign_add(tf$reduce_sum(inputs, axis=0L))
    self$total
  }
)

my_sum <- ComputeSum(input_dim = 2)
x <- tf$ones(shape(2, 2))

as.array(my_sum(x))

as.array(my_sum(x))

my_sum$trainable_weights

# Let's reuse the Linear class
# with a `build` method that we defined above.

MLP <- new_layer_class(
  "MLP",
  initialize = function() {
    super$initialize()
    self$linear_1 <- Linear(units = 32)
    self$linear_2 <- Linear(units = 32)
    self$linear_3 <- Linear(units = 10)
  },
  call = function(inputs) {
    x <- self$linear_1(inputs)
    x <- tf$nn$relu(x)
    x <- self$linear_2(x)
    x <- tf$nn$relu(x)
    return(self$linear_3(x))
  }
)

mlp <- MLP()

# The first call to the `mlp` object will create the weights.
y <- mlp(tf$ones(shape=shape(3, 64)))

# Weights are recursively tracked.
length(mlp$weights)

mlp <- keras_model_sequential() %>%
  layer_dense(units = 32, activation = "relu") %>%
  layer_dense(units = 32, activation = "relu") %>%
  layer_dense(units = 10)

# A layer that creates an activity sparsity regularization loss
ActivityRegularization <- new_layer_class(
  "ActivityRegularization",
  initialize = function(rate=1e-2) {
    super$initialize()
    self$rate <- rate
  },
  call = function(inputs) {
    self$add_loss(self$rate * tf$reduce_sum(tf$abs(inputs)))
    inputs
  }
)

# Let's use the loss layer in a MLP block.
SparseMLP <- new_layer_class(
  "SparseMLP",
  initialize = function() {
    super$initialize()
    self$linear_1 <- Linear(units = 32)
    self$reg <- ActivityRegularization(rate = 1e-2)
    self$linear_3 <- Linear(units = 10)
  },
  call = function(inputs) {
    x <- self$linear_1(inputs)
    x <- tf$nn$relu(x)
    x <- self$reg(x)
    return(self$linear_3(x))
  }
)

mlp <- SparseMLP()
y <- mlp(tf$ones(shape(10, 10)))

mlp$losses  # List containing one float32 scalar

# Losses correspond to the *last* forward pass.
mlp <- SparseMLP()
mlp(tf$ones(shape(10, 10)))

length(mlp$losses)

mlp(tf$ones(shape(10, 10)))

length(mlp$losses)  # No accumulation.

# Let's demonstrate how to use these losses in a training loop.

# Prepare a dataset.
c(c(x_train, y_train), .) %<-% dataset_mnist()
x_train <- array_reshape(x_train, c(60000, 784)) / 255

dataset <- tfdatasets::tensor_slices_dataset(list(x_train, y_train)) %>%
  tfdatasets::dataset_shuffle(buffer_size=1024) %>%
  tfdatasets::dataset_batch(64)

# A new MLP.
mlp <- SparseMLP()

# Loss and optimizer.
loss_fn <- loss_sparse_categorical_crossentropy(from_logits=TRUE)
optimizer <- optimizer_sgd(learning_rate=1e-3)

coro::loop(for(data in dataset) {
  x <- data[[1]]
  y <- data[[2]]
  with(tf$GradientTape() %as% tape, {
    # Forward pass.
    logits <- mlp(x)

    # External loss value for this batch.
    loss <- loss_fn(y, logits)

    # Add the losses created during the forward pass.
    loss <- loss + Reduce(`+`, mlp$losses)

    # Get gradients of the loss wrt the weights.
    gradients <- tape$gradient(loss, mlp$trainable_weights)

    # Update the weights of our linear layer.
    optimizer$apply_gradients(zip_lists(gradients, mlp$trainable_weights))
  })
})

# Instantiate a metric object
accuracy <- metric_sparse_categorical_accuracy()

# Prepare our layer, loss, and optimizer.
model <- keras_model_sequential() %>%
  layer_dense(units = 32, activation = "relu") %>%
  layer_dense(units = 32, activation = "relu") %>%
  layer_dense(units = 10)
loss_fn <- loss_sparse_categorical_crossentropy(from_logits = TRUE)
optimizer <- optimizer_adam(learning_rate=1e-3)

for (epoch in seq_len(2)) {
  coro::loop(for (data in dataset) {
    x <- data[[1]]
    y <- data[[2]]
    with(tf$GradientTape() %as% tape, {
      # Forward pass.
      logits <- model(x)

      # External loss value for this batch.
      loss_value <- loss_fn(y, logits)
    })

    # Update the state of the `accuracy` metric.
    accuracy$update_state(y, logits)

    # Update the weights of the model to minimize the loss value.
    gradients <- tape$gradient(loss_value, model$trainable_weights)
    optimizer$apply_gradients(zip_lists(gradients, model$trainable_weights))

  })
  cat("Epoch:", epoch, "Accuracy:", as.numeric(accuracy$result()), "\n")
  accuracy$reset_state()
}

F1Score <- new_metric_class(
  "F1Score",
  initialize = function(self, name="f1_score", dtype="float32", threshold=0.5, ...) {
    super$initialize(name=name, dtype=dtype, ...)
    self$threshold <- threshold
    self$true_positives <- self$add_weight(
      name="tp", dtype=dtype, initializer="zeros"
    )
    self$false_positives <- self$add_weight(
      name="fp", dtype=dtype, initializer="zeros"
    )
    self$false_negatives <- self$add_weight(
      name="fn", dtype=dtype, initializer="zeros"
    )
  },
  update_state = function(y_true, y_pred, sample_weight=NULL) {
    y_pred <- tf$math$greater_equal(y_pred, self$threshold)
    y_true <- tf$cast(y_true, tf$bool)
    y_pred <- tf$cast(y_pred, tf$bool)

    true_positives <- tf$cast(y_true & y_pred, self$dtype)
    false_positives <- tf$cast((!y_true) & y_pred, self$dtype)
    false_negatives <- tf$cast(y_true & (!y_pred), self$dtype)

    if (!is.null(sample_weight)) {
      sample_weight <- tf$cast(sample_weight, self$dtype)
      true_positives <- true_positives * sample_weight
      false_positives <- false_positives * sample_weight
      false_negatives <- false_negatives * sample_weight
    }

    self$true_positives$assign_add(tf$reduce_sum(true_positives))
    self$false_positives$assign_add(tf$reduce_sum(false_positives))
    self$false_negatives$assign_add(tf$reduce_sum(false_negatives))
  },

  result = function() {
    precision <- self$true_positives / (self$true_positives + self$false_positives)
    recall <- self$true_positives / (self$true_positives + self$false_negatives)
    f1_score <- 2 * precision * recall / (precision + recall)
    f1_score
  },

  reset_state = function() {
    self$true_positives$assign(0)
    self$false_positives$assign(0)
    self$false_negatives$assign(0)
  }
)

m <- F1Score()
m$update_state(c(0, 1, 0, 0), c(0.3, 0.5, 0.8, 0.9))
cat("Intermediate result:", as.numeric(m$result()), "\n")

m$update_state(c(1, 1, 1, 1), c(0.1, 0.7, 0.6, 0.0))
cat("Final result:", as.numeric(m$result()), "\n")

# Prepare our layer, loss, and optimizer.
model <- keras_model_sequential() %>%
  layer_dense(units = 32, activation = "relu") %>%
  layer_dense(units = 32, activation = "relu") %>%
  layer_dense(units = 10)
loss_fn <- loss_sparse_categorical_crossentropy(from_logits = TRUE)
optimizer <- optimizer_adam(learning_rate=1e-3)

# Create a training step function.
train_on_batch <- tf_function(function(x, y) {
  with(tf$GradientTape() %as% tape, {
    # Forward pass.
    logits <- model(x)
    # External loss value for this batch.
    loss_value <- loss_fn(y, logits)
  })
  # Update the weights of the model to minimize the loss value.
  gradients <- tape$gradient(loss_value, model$trainable_weights)
  optimizer$apply_gradients(zip_lists(gradients, model$trainable_weights))
  loss_value
})


# Prepare a dataset.
c(c(x_train, y_train), .) %<-% dataset_mnist()
x_train <- array_reshape(x_train, c(60000, 784)) / 255

dataset <- tfdatasets::tensor_slices_dataset(list(x_train, y_train)) %>%
  tfdatasets::dataset_shuffle(buffer_size=1024) %>%
  tfdatasets::dataset_batch(64)

i <- 0
coro::loop(for (data in dataset) {
  i <- i + 1
  x <- data[[1]]
  y <- data[[2]]
  loss <- train_on_batch(x, y)
  if (i %% 100 == 0)
    cat("Loss:", as.numeric(loss), "\n")
})

Dropout <- new_layer_class(
  "Dropout",
  initialize = function(rate) {
    super$initialize()
    self$rate <- rate
  },
  call = function(inputs, training = NULL) {
    if (!is.null(training) && training) {
      return(tf$nn$dropout(inputs, rate = self$rate))
    }
    inputs
  }
)

MLPWithDropout <- new_layer_class(
  "MLPWithDropout",
  initialize = function() {
    super$initialize()
    self$linear_1 <- Linear(units = 32)
    self$dropout <- Dropout(rate = 0.5)
    self$linear_3 <- Linear(units = 10)
  },
  call = function(inputs, training = NULL) {
    x <- self$linear_1(inputs)
    x <- tf$nn$relu(x)
    x <- self$dropout(x, training = training)
    self$linear_3(x)
  }
)

mlp <- MLPWithDropout()
y_train <- mlp(tf$ones(shape(2, 2)), training=TRUE)
y_test <- mlp(tf$ones(shape(2, 2)), training=FALSE)

# We use an `Input` object to describe the shape and dtype of the inputs.
# This is the deep learning equivalent of *declaring a type*.
# The shape argument is per-sample; it does not include the batch size.
# The functional API focused on defining per-sample transformations.
# The model we create will automatically batch the per-sample transformations,
# so that it can be called on batches of data.
inputs <- layer_input(shape = 16, dtype = "float32")

# We call layers on these "type" objects
# and they return updated types (new shapes/dtypes).
outputs <- inputs %>%
  Linear(units = 32) %>% # We are reusing the Linear layer we defined earlier.
  Dropout(rate = 0.5) %>% # We are reusing the Dropout layer we defined earlier.
  Linear(units = 10)

# A functional `Model` can be defined by specifying inputs and outputs.
# A model is itself a layer like any other.
model <- keras_model(inputs, outputs)

# A functional model already has weights, before being called on any data.
# That's because we defined its input shape in advance (in `Input`).
length(model$weights)

# Let's call our model on some data, for fun.
y <- model(tf$ones(shape(2, 16)))
y$shape

# You can pass a `training` argument in `__call__`
# (it will get passed down to the Dropout layer).
y <- model(tf$ones(shape(2, 16)), training=TRUE)

inputs <- layer_input(shape = 784, dtype="float32")
outputs <- inputs %>%
  layer_dense(units = 32, activation = "relu") %>%
  layer_dense(units = 32, activation = "relu") %>%
  layer_dense(units = 10)
model <- keras_model(inputs, outputs)

# Specify the loss, optimizer, and metrics with `compile()`.
model %>% compile(
    loss = loss_sparse_categorical_crossentropy(from_logits=TRUE),
    optimizer=optimizer_adam(learning_rate=1e-3),
    metrics=list(metric_sparse_categorical_accuracy()),
)

# Train the model with the dataset for 2 epochs.
model %>% fit(dataset, epochs=2)

predictions <- model %>% predict(dataset)

model %>% evaluate(dataset)

CustomModel <- new_model_class(
  "CustomModel",
  initialize = function(...) {
    super$initialize(...)
    self$loss_tracker <- metric_mean(name="loss")
    self$accuracy <- metric_sparse_categorical_accuracy()
    self$loss_fn <- loss_sparse_categorical_crossentropy(from_logits=TRUE)
    self$optimizer <- optimizer_adam(learning_rate=1e-3)
  },
  train_step = function(data) {
    c(x, y, sample_weight) %<-% unpack_x_y_sample_weight(data)
    with(tf$GradientTape() %as% tape, {
      y_pred <- self(x, training=TRUE)
      loss <- self$loss_fn(y = y, y_pred = y_pred, sample_weight=sample_weight)
    })
    gradients <- tape$gradient(loss, self$trainable_variables)
    self$optimizer$apply_gradients(
      zip_lists(gradients, self$trainable_variables)
    )

    # Update metrics (includes the metric that tracks the loss)
    self$loss_tracker$update_state(loss)
    self$accuracy$update_state(y, y_pred, sample_weight=sample_weight)
    # Return a list mapping metric names to current value
    list(
      loss = self$loss_tracker$result(),
      accuracy = self$accuracy$result()
    )
  },
  metrics = mark_active(function() {
    list(self$loss_tracker, self$accuracy)
  })
)

inputs <- layer_input(shape = 784, dtype="float32")
outputs <- inputs %>%
  layer_dense(units = 32, activation = "relu") %>%
  layer_dense(units = 32, activation = "relu") %>%
  layer_dense(units = 10)
model <- CustomModel(inputs, outputs)
model %>% compile()
model %>% fit(dataset, epochs=2)

Sampling <- new_layer_class(
  "Sampling",
  call = function(inputs) {
    c(z_mean, z_log_var) %<-% inputs
    batch <- op_shape(z_mean)[[1]]
    dim <- op_shape(z_mean)[[2]]
    epsilon <- random_normal(shape = c(batch, dim))
    z_mean + op_exp(0.5 * z_log_var) * epsilon
  }
)

Encoder <- new_layer_class(
  "Encoder",
  initialize = function(latent_dim = 32, intermediate_dim = 64, ...) {
    super$initialize(...)
    self$dense_proj <- layer_dense(units = intermediate_dim, activation = "relu")
    self$dense_mean <- layer_dense(units = latent_dim)
    self$dense_log_var <- layer_dense(units = latent_dim)
    self$sampling <- Sampling()
  },
  call = function(inputs) {
    x <- self$dense_proj(inputs)
    z_mean <- self$dense_mean(x)
    z_log_var <- self$dense_log_var(x)
    z <- self$sampling(list(z_mean, z_log_var))
    list(z_mean, z_log_var, z)
  }
)

Decoder <- new_layer_class(
  "Decoder",
  initialize = function(original_dim, intermediate_dim = 64, ...) {
    super$initialize(...)
    self$dense_proj <- layer_dense(units = intermediate_dim, activation = "relu")
    self$dense_output <- layer_dense(units = original_dim, activation = "sigmoid")
  },
  call = function(inputs) {
    x <- self$dense_proj(inputs)
    self$dense_output(x)
  }
)

VariationalAutoEncoder <- new_model_class(
  "VariationalAutoEncoder",
  initialize = function(original_dim,
        intermediate_dim=64,
        latent_dim=32,
        name="autoencoder", ...) {
    super$initialize(name = name, ...)
    self$original_dim <- original_dim
    self$encoder <- Encoder(
      latent_dim = latent_dim,
      intermediate_dim = intermediate_dim
    )
    self$decoder <- Decoder(
      original_dim = original_dim,
      intermediate_dim = intermediate_dim
    )
  },
  call = function(inputs) {
    c(z_mean, z_log_var, z) %<-% self$encoder(inputs)
    reconstructed <- self$decoder(z)
    # Add KL divergence regularization loss.
    kl_loss <- -0.5 * op_mean(
      z_log_var - op_square(z_mean) - op_exp(z_log_var) + 1
    )
    self$add_loss(kl_loss)
    reconstructed
  }
)

# Our model.
vae <- VariationalAutoEncoder(
  original_dim = 784,
  intermediate_dim = 64,
  latent_dim = 32
)

# Loss and optimizer.
loss_fn <- loss_mean_squared_error()
optimizer = optimizer_adam(learning_rate=1e-3)

# Prepare a dataset.
c(c(x_train, .), .) %<-% dataset_mnist()
x_train <- array_reshape(x_train, c(60000, 784)) / 255

dataset <- tfdatasets::tensor_slices_dataset(x_train) %>%
  tfdatasets::dataset_shuffle(buffer_size=1024) %>%
  tfdatasets::dataset_batch(32)


training_step <- tf_function(function(x) {
  with(tf$GradientTape() %as% tape, {
    reconstructed <- vae(x)  # Compute input reconstruction.
    # Compute loss.
    loss <- loss_fn(x, reconstructed)
    loss <- loss + op_sum(vae$losses)  # Add KLD term.
  })
  # Update the weights of the VAE.
  grads <- tape$gradient(loss, vae$trainable_weights)
  optimizer$apply_gradients(zip_lists(grads, vae$trainable_weights))
  loss
})

losses <- c()  # Keep track of the losses over time.
coro::loop(for(data in dataset) {
  loss <- training_step(data)

  # Logging.
  losses[length(losses) + 1] <- as.numeric(loss)
  if (length(losses) %% 100 == 0) {
    cat("Step:", length(losses), "Loss:", mean(losses), "\n")
  }
  # Stop after 1000 steps.
  # Training the model to convergence is left
  # as an exercise to the reader.
  if (length(losses) >= 1000) {
    break
  }
})

input_dim <- 784
classes <- 10

# This is the main network we'll actually use to predict labels.
inputs <- layer_input(shape = input_dim)
dense1 <- layer_dense(units = 64, activation = "relu")
dense1$built <- TRUE

dense2 <- layer_dense(units = classes)
dense2$built <- TRUE

outputs <- inputs %>% dense1() %>% dense2()
main_network <- keras_model(inputs, outputs)

# This is the number of weight coefficients to generate. Each layer in the
# main network requires output_dim * input_dim + output_dim coefficients.
num_weights_to_generate <- (classes * 64 + classes) + (64 * input_dim + 64)

# This is the hypernetwork that generates the weights of the `main_network` above.
hypernetwork <- keras_model_sequential() %>%
  layer_dense(units=16, activation="relu") %>%
  layer_dense(units=num_weights_to_generate, activation="sigmoid")

# Loss and optimizer.
loss_fn <- loss_sparse_categorical_crossentropy(from_logits = TRUE)
optimizer <- optimizer_adam(learning_rate=1e-4)

# Prepare a dataset.
c(c(x_train, y_train), .) %<-% dataset_mnist()
x_train <- array_reshape(x_train, c(60000, 784)) / 255

dataset <- tfdatasets::tensor_slices_dataset(list(x_train, y_train)) %>%
  tfdatasets::dataset_shuffle(buffer_size=1024) %>%
  # We'll use a batch size of 1 for this experiment.
  tfdatasets::dataset_batch(1)

train_step <- function(x, y) {
  with(tf$GradientTape() %as% tape, {
    weights_pred <- hypernetwork(x)

    # Reshape them to the expected shapes for w and b for the outer model.
    # Layer 1 kernel.
    start_index <- 1
    w1_shape <- c(input_dim, 64)
    w1_coeffs <- weights_pred[, start_index:(start_index + prod(w1_shape) - 1)]
    w1 <- tf$reshape(w1_coeffs, as.integer(w1_shape))
    start_index <- start_index + prod(w1_shape)

    # Layer 1 bias.
    b1_shape <- c(64)
    b1_coeffs <- weights_pred[, start_index:(start_index + prod(b1_shape) - 1)]
    b1 <- tf$reshape(b1_coeffs, as.integer(b1_shape))
    start_index <- start_index + prod(b1_shape)

    # Layer 2 kernel.
    w2_shape <- c(64, classes)
    w2_coeffs <- weights_pred[, start_index:(start_index + prod(w2_shape) - 1)]
    w2 <- tf$reshape(w2_coeffs, as.integer(w2_shape))
    start_index <- start_index + prod(w2_shape)

    # Layer 2 bias.
    b2_shape <- c(classes)
    b2_coeffs <- weights_pred[, start_index:(start_index + prod(b2_shape) - 1)]
    b2 <- tf$reshape(b2_coeffs, as.integer(b2_shape))
    start_index <- start_index + prod(b2_shape)

    # Set the weight predictions as the weight variables on the outer model.
    dense1$kernel <- w1
    dense1$bias <- b1
    dense2$kernel <- w2
    dense2$bias <- b2

    # Inference on the outer model.
    preds <- main_network(x)
    loss <- loss_fn(y, preds)
  })

  grads <- tape$gradient(loss, hypernetwork$trainable_weights)
  optimizer$apply_gradients(zip_lists(grads, hypernetwork$trainable_weights))
  loss
}

losses <- c()  # Keep track of the losses over time.
coro::loop(for (data in dataset) {
  x <- data[[1]]
  y <- data[[2]]
  loss <- train_step(x, y)

  # Logging.
  losses[length(losses) + 1] <- as.numeric(loss)
  if (length(losses) %% 100 == 0) {
    cat("Step:", length(losses), "Loss:", mean(losses), "\n")
  }
  # Stop after 1000 steps.
  # Training the model to convergence is left
  # as an exercise to the reader.
  if (length(losses) >= 1000) {
    break
  }
})