import tensorflow as tf

        cross_entropy = tf.reduce_mean(
            -tf.reduce_sum(y_ * tf.log(y), reduction_indices=[1])
        )
        self.cross_entropy = cross_entropy
        self.cross_entropy_grads = tf.gradients(cross_entropy, [w, b])
        self.sess = tf.Session()
        # In order to get and set the weights, we pass in the loss function to
        # Ray's TensorFlowVariables to automatically create methods to modify
        # the weights.
        self.variables = ray.experimental.tf_utils.TensorFlowVariables(
