message: Correct benchmark parameters for inputs.
diff: @@ -162,15 +162,15 @@ class KerasLayerBenchmarksBackwardXLA(six.with_metaclass(
  _benchmark_parameters = _generate_benchmark_params([
  ("Conv2D_small_shape", tf.keras.layers.Conv2D,
  {"filters": 1, "kernel_size": 1, "activation": "relu"},
  - (1, 1, 1, 1), 10000),
  + {"input_shape": (1, 1, 1, 1)}, 10000),
  ("Conv2D_normal_shape", tf.keras.layers.Conv2D,
  {"filters": 1, "kernel_size": 1, "activation": "relu"},
  - (64, 28, 28, 3), 10000),
  + {"input_shape": (64, 28, 28, 3)}, 10000),
  # TODO(b/153480400)
  # ("LSTM_small_shape", tf.keras.layers.LSTM,
  - # {"units": 1}, (1, 1, 1), 10000),
  + # {"units": 1}, {"input_shape": (1, 1, 1)}, 10000),
  # ("LSTM_normal_shape", tf.keras.layers.LSTM,
  - # {"units": 4}, (32, 10, 8), 10000),
  + # {"units": 4}, {"input_shape": (32, 10, 8)}, 10000),
  ("Embedding_small_shape", tf.keras.layers.Embedding,
  {"input_dim": 1, "output_dim": 1, "input_length": 1},
  {"input": np.random.randint(1, size=(1, 1))}, 10),
  