code: def inputwise_condition_numbers(f, *args):
      """Computes the condition numbers of `f(*args)` at each arg independently.
  
      The function `f(*args)` must produce a scalar result; computing
      batches of condition numbers or computing condition numbers of
      vector-valued functions is not yet supported.
  
      This function assumes that numerical error when computing `f` in
      float64 is negligible.  For this to work correctly, `f` needs to be
      _dtype-polymorphic_: the dtype in which computations internal to `f`
      are performed should match the dtype of the arguments of `f`.
  
      Args:
        f: Function whose accuracy to evaluate.  Must be differentiable
          and dtype-polymorphic.
        *args: Arguments at which to test the accuracy of `f`.
  
      Returns:
        condition_numbers: The condition number of `f` with respect to each input.
          The returned structure is parallel to `*args`.
  
      Raises:
        ValueError: If `f` is found not to be dtype-polymorphic.
  
      """
      # TODO(b/181967692): Compute multivariate condition numbers.
      # TODO(b/181967437): To support batch condition numbers, need batch gradients.
      # Then can infer the "event shape" of the arguments by subtracting
      # off the number of dimensions in f(*args).
      # To also support vector outputs, need to know the "event_ndims" in
      # the output f(*args), and need full Jacobians of f underneath.
      args_32 = tf.nest.map_structure(floating_tensor_to_f32, args)
      logging.vlog(1, "32-bit arguments: %s", args_32)
      args_64 = tf.nest.map_structure(floating_tensor_to_f64, args_32)
      truth, derivatives = gradient.value_and_gradient(f, args_64)
      logging.vlog(1, "Correct answer: %s", truth)
      logging.vlog(1, "Argument gradient: %s", derivatives)
  
      def check_numerics(x):
          if x is None:
              return None
          msg = "Cannot check accuracy if ground truth or derivatives are not finite"
          return tf.debugging.check_numerics(x, message=msg)
  
      truth = check_numerics(truth)
      derivatives = tf.nest.map_structure(check_numerics, derivatives)
      if truth.dtype != tf.float64:
          raise ValueError(
              "Evaluating on {} produced non-64-bit result {}".format(
                  args_64, truth
              )
          )
      return tf.nest.map_structure(
          functools.partial(condition_number_one_input, truth),
          # For some reason, value_and_gradient casts the outer structure to list in
          # jax.  Is that an oversight?
          tuple(args_64),
          tuple(derivatives),
      )
  
file_path: probability/tensorflow_probability/python/internal/numerics_testing.py
imports: import functools
  from absl import logging
  import numpy
  import tensorflow.compat.v2
  from tensorflow_probability.python.internal import dtype_util
  from tensorflow_probability.python.math import gradient
label_desc: Use image node_id 9 for calling a global function with example usage: inputwise_condition_numbers(f) and returns: tf, tf, None