def _get_reduce_dims(op):
  """Returns the reduction dimensions for grouping weights of various ops."""
  type_to_dims = {'Conv2D': (0, 1, 2), 'Conv2DBackpropInput': (0, 1, 3)}
  try:
    return type_to_dims[op.type]
  except KeyError:
    raise ValueError('Reduce dims are unknown for op type %s' % op.type)