def __init__(self, canvas_shape,
               num_batch_axes=1,
               dtype=tf.float32,
               **kwargs):
    super(Cache, self).__init__(trainable=False, **kwargs)
    self.canvas_shape = canvas_shape
    self.num_batch_axes = num_batch_axes
    self._dtype = dtype