async def find_failing_ops_with_infinite_output(
        self,
        inputs: Dict[str, np.array],
        output_names: List[str] = None,
        skip_op: Callable[[proto.MIL_pb2.Operation], bool] = skip_op_by_type,
    ) -> List[proto.MIL_pb2.Operation]:
        """
        Identify operations in the model that produce infinite outputs.

        This method traverses the model's operation graph and checks for infinite values
        in the output of each operation. It returns a list of operations that are likely the
        source of failures. It's useful for debugging numerical instability issues in the model.

        Parameters
        ----------
        inputs : dict[str, np.array]
            A dictionary of input tensors for the model. Keys are input names, and values are numpy arrays.

        output_names : Optional[List[str]]
            A list of specific output names to start the traversal from. If None, all model outputs are used.

        skip_op: Callable[[proto.MIL_pb2.Operation], bool]
            A function that determines if an operation should be skipped.

        Returns
        -------
        List[proto.MIL_pb2.Operation]
            A list of operations that are identified as likely sources of failure.

        Notes
        -----
            - The method uses a breadth-first search strategy to traverse the operation graph.
            - An operation is considered a failure source if it outputs infinite values while its direct inputs do not.

        Examples
        --------
        .. sourcecode:: python

            validator = coremltools.models.ml_program.experimental.debugging_utils.MLModelValidator(
                model
            )
            failing_ops = await validator.find_failing_ops_with_infinite_output(
                inputs={"input": np.array([1, 2, 3])}, output_names=["output1", "output2"]
            )
        """

        def validate_output(op: proto.MIL_pb2.Operation, value: np.array):
            return not np.isinf(value).any()

        return await self.find_failing_ops(
            inputs=inputs,
            validate_output=validate_output,
            output_names=output_names,
            skip_op=skip_op,
        )

        Identify operations in the model that produce infinite outputs.

        This method traverses the model's operation graph and checks for infinite values
        in the output of each operation. It returns a list of operations that are likely the
        source of failures. It's useful for debugging numerical instability issues in the model.

        Parameters
        ----------
        inputs : dict[str, np.array]
            A dictionary of input tensors for the model. Keys are input names, and values are numpy arrays.

        output_names : Optional[List[str]]
            A list of specific output names to start the traversal from. If None, all model outputs are used.

        skip_op: Callable[[proto.MIL_pb2.Operation], bool]
            A function that determines if an operation should be skipped.

        Returns
        -------
        List[proto.MIL_pb2.Operation]
            A list of operations that are identified as likely sources of failure.

        Notes
        -----
            - The method uses a breadth-first search strategy to traverse the operation graph.
            - An operation is considered a failure source if it outputs infinite values while its direct inputs do not.

        Examples
        --------
        .. sourcecode:: python

            validator = coremltools.models.ml_program.experimental.debugging_utils.MLModelValidator(
                model
            )
            failing_ops = await validator.find_failing_ops_with_infinite_output(
                inputs={"input": np.array([1, 2, 3])}, output_names=["output1", "output2"]
            )
        
coremltools/models/ml_program/experimental/debugging_utils.py