def check_unique_axis_names_for_initializer(
    input_axes: dict[str, int],
    output_axes: dict[str, int],
    parallel_axes: dict[str, int],
    convolution_spatial_axes: dict[str, int],
) -> None:
  """Checks that indices appear exactly once in the given axis specs."""
  seen_counts = collections.Counter(
      itertools.chain(
          input_axes.keys(),
          output_axes.keys(),
          parallel_axes.keys(),
          convolution_spatial_axes.keys(),
      )
  )
  bad_names = {name for name, count in seen_counts.items() if count > 1}
  if bad_names:
    raise ValueError(
        f"Duplicate axis names during initialization: {bad_names}."
    )
Checks that indices appear exactly once in the given axis specs.
penzai/nn/linear_and_affine.py