Custom metric with 2d weights is not working

[RazvanPasca]

System information

• Have I written custom code (as opposed to using a stock example script
  provided in TensorFlow Model Analysis): yes
• OS Platform and Distribution (e.g., Linux Ubuntu 16.04): Linux Ubuntu 20.04
• TensorFlow Model Analysis version (use command below): 0.29.0
• Tensorflow version: 2.4.1
• TFX version: 0.29.0
• Python version: 3.8

Describe the problem

I have implemented a Personalization metric as a tf.keras.Metric subclass which is used as input for an Evaluator component for a local TFX pipeline. The evaluation step crashes because of some input size mismatch when running the add_weights

model-analysis/tensorflow_model_analysis/metrics/tf_metric_accumulators.py

Line 110 in cc7d75c

def add_weights(self, output_index: int, metric_index: int,

method in the _CompilableMetricsCombiner class. For other custom metrics that I implemented which compute averages over batches (hence scalars, not 2d) it does not happen.

The training step of the TFX pipeline works correctly and manages to run the metric without problems, hence I think it is not a problem of implementation, but of how tfma handles/wraps custom keras metrics.

I tried to hack it by changing the value of TFMetricsAccumulator._DEFAULT_DESIRED_BATCH_SIZE to 1 and 2 to avoid size mismatches but then the values are completely off (probably some value aggregations do not like it)

Source code / logs

The code for the metric implementation is found here

The traceback for the error is found here
Last lines:
File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tensorflow_model_analysis/evaluators/metrics_plots_and_validations_evaluator.py", line 396, in compact return super(_ComputationsCombineFn, self).compact(accumulator) File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/transforms/combiners.py", line 756, in compact return [ File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/apache_beam/transforms/combiners.py", line 757, in <listcomp> c.compact(a, *args, **kwargs) for c, File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tensorflow_model_analysis/metrics/tf_metric_wrapper.py", line 568, in compact self._process_batch(accumulator) File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tensorflow_model_analysis/metrics/tf_metric_wrapper.py", line 497, in _process_batch accumulator.add_weights(output_index, metric_index, File "/home/user/.local/share/virtualenvs/project-xv7hiZ8W/lib/python3.8/site-packages/tensorflow_model_analysis/metrics/tf_metric_accumulators.py", line 120, in add_weights self._weights[output_index][metric_index] = np.add(cur_weights, weights) ValueError: operands could not be broadcast together with shapes (1,937,30208) (1,73,30208)

[RazvanPasca]

Hello! Any updates on this?