Wrap forward and backward pass in a defun for L2HMC.

Also a small bugfix to handle unknown shapes in backprop._num_elements.

Before:
entry {
  name: "L2hmcBenchmark.eager_train_cpu_defun"
  iters: 10
  wall_time: 0.594115018845
  extras {
    key: "examples_per_sec"
    value {
      double_value: 336.635152548
    }
  }
}

After:
entry {
  name: "L2hmcBenchmark.eager_train_cpu_defun"
  iters: 10
  wall_time: 0.322251081467
  extras {
    key: "examples_per_sec"
    value {
      double_value: 620.634069216
    }
  }
}
PiperOrigin-RevId: 214308142