Add input_shape metric to jagged_sum operator (#2313)

Summary:
Pull Request resolved: #2313

Add new metric to `jagged_sum` that denotes the 0th and 2nd input dimensions, `B` and `M`, in the form `(B, '*', M)`, where the nested tensor has logical dimensions `(B, *, M)`. Display this metric once per `x` value using the `x_only = True` argument to `register_metric()`. This diff will make TritonBench's benchmark table more readable by denoting the nested tensor dimensions used per benchmark row.

Reviewed By: jbschlosser

Differential Revision: D58535619

fbshipit-source-id: dcc5f4efab0b81d605689f7371edee7ddc511caf

torchbenchmark/operators/jagged_sum/operator.py

@@ -156,3 +156,13 @@ def gbps(self, fn_name, example_inputs, metrics: BenchmarkOperatorMetrics):
  / metrics.latency
  * GIGABYTES_PER_BYTE
  )
+
+ @register_metric(x_only=True)
+ def input_shape(
+ self, fn_name: str, example_inputs, metrics: BenchmarkOperatorMetrics
+ ):
+ return (
+ example_inputs[0].shape[0],
+ "*",
+ example_inputs[0].shape[2],
+ ) # return (B, '*', M) for each example input