Some tiny speed improvements for defun.

Before:
entry {
  name: "MicroBenchmarks.benchmark_defun_matmul_2_by_2_CPU"
  iters: 30000
  wall_time: 48.4476327896
  extras {
    key: "examples_per_sec"
    value {
      double_value: 20640.8433688
    }
  }
}

After:
entry {
  name: "MicroBenchmarks.benchmark_defun_matmul_2_by_2_CPU"
  iters: 30000
  wall_time: 45.2344338099
  extras {
    key: "examples_per_sec"
    value {
      double_value: 22107.0524327
    }
  }
}
PiperOrigin-RevId: 215619902