import re


def count_ops_in_hlo_proto(hlo_proto,
                           ops_regex):
  """Counts specific ops in hlo proto, whose names match the provided pattern.

  Args:
    hlo_proto: an HloModuleProto object.
    ops_regex: a string regex to filter ops with matching names.

  Returns:
    Count of matching ops.

  """
  ops_count = 0
  for computation in hlo_proto.computations:
    for instr in computation.instructions:
      if re.match(ops_regex, instr.name) is not None:
        ops_count += 1
  return ops_count