Minor eager mode improvements

Mainly doesn't get shape when getting rank will suffice.

Before:
entry {
  name: "EagerLinearRegressionBenchmark.eager_train_cpu"
  iters: 2000
  wall_time: 1.54617094994
  extras {
    key: "examples_per_sec"
    value {
      double_value: 82785.1538701
    }
  }
}

After:
entry {
  name: "EagerLinearRegressionBenchmark.eager_train_cpu"
  iters: 2000
  wall_time: 1.38547611237
  extras {
    key: "examples_per_sec"
    value {
      double_value: 92387.0132856
    }
  }
}
PiperOrigin-RevId: 204001196