Compute non-zero flops for gather with FILL_OR_DROP.

[copybara-service]

Compute non-zero flops for gather with FILL_OR_DROP.

Example lowered HLO for the FILL_OR_DROP test (14 flops):

module @jit__lambda attributes {mhlo.num_partitions = 1 : i32, mhlo.num_replicas = 1 : i32} {
  func.func public @main(%arg0: tensor<3x3xi32>, %arg1: tensor<2x1xi32>) -> (tensor<2x3xi32> {jax.result_info = "result"}) {
    %c = stablehlo.constant dense<2> : tensor<1xi64>
    %0 = stablehlo.convert %arg1 : (tensor<2x1xi32>) -> tensor<2x1xi64>
    %c_0 = stablehlo.constant dense<0> : tensor<i64>
    %1 = stablehlo.broadcast_in_dim %c_0, dims = [] : (tensor<i64>) -> tensor<2x1xi64>
    %2 = stablehlo.compare  GE, %0, %1,  SIGNED : (tensor<2x1xi64>, tensor<2x1xi64>) -> tensor<2x1xi1>
    %3 = stablehlo.broadcast_in_dim %c, dims = [1] : (tensor<1xi64>) -> tensor<1x1xi64>
    %4 = stablehlo.broadcast_in_dim %3, dims = [0, 1] : (tensor<1x1xi64>) -> tensor<2x1xi64>
    %5 = stablehlo.compare  LE, %0, %4,  SIGNED : (tensor<2x1xi64>, tensor<2x1xi64>) -> tensor<2x1xi1>
    %6 = stablehlo.and %2, %5 : tensor<2x1xi1>
    %c_1 = stablehlo.constant dense<true> : tensor<i1>
    %7 = stablehlo.reduce(%6 init: %c_1) applies stablehlo.and across dimensions = [1] : (tensor<2x1xi1>, tensor<i1>) -> tensor<2xi1>
    %8 = "stablehlo.gather"(%arg0, %0) <{dimension_numbers = #stablehlo.gather<offset_dims = [1], collapsed_slice_dims = [0], start_index_map = [0], index_vector_dim = 1>, indices_are_sorted = false, slice_sizes = array<i64: 1, 3>}> : (tensor<3x3xi32>, tensor<2x1xi64>) -> tensor<2x3xi32>
    %9 = stablehlo.broadcast_in_dim %7, dims = [0] : (tensor<2xi1>) -> tensor<2x3xi1>
    %c_2 = stablehlo.constant dense<-2147483648> : tensor<i32>
    %10 = stablehlo.broadcast_in_dim %c_2, dims = [] : (tensor<i32>) -> tensor<2x3xi32>
    %11 = stablehlo.select %9, %8, %10 : tensor<2x3xi1>, tensor<2x3xi32>
    return %11 : tensor<2x3xi32>
  }
}