def _fan_in_concat_shape(input_shape, axis=-1):  # pylint: disable=invalid-name
  """Helper to determine the shape of FanInConcat output."""
  ax = axis % len(input_shape[0])
  concat_size = sum(shape[ax] for shape in input_shape)
  out_shape = input_shape[0][:ax] + (concat_size,) + input_shape[0][ax+1:]
  return out_shape