guides/writing_a_custom_training_loop_in_jax.py/ActivityRegularizationLayer/call
class ActivityRegularizationLayer:    def call(self, inputs):
        self.add_loss(1e-2 * jax.numpy.sum(inputs))
        return inputs