(jacobian_array)
return {
      "metric_tensor": _analyze_metric_tensor(metric_tensor),
      "mean_metric_tensor": _analyze_metric_tensor(mean_metric_tensor)
  }
Computes eigenvalue statistics of the Jacobian.

  Computes the eigenvalues and condition number of the metric tensor for the
  Jacobian evaluated at each element of the batch and the mean metric tensor
  across the batch.

  Args:
    jacobian_array: A numpy array holding the Jacobian.

  Returns:
    A dict of spectral statistics with two elements, one containing stats
    for every metric tensor in the batch, another for the mean metric tensor.
Computes eigenvalue statistics of the Jacobian.
Computes eigenvalue statistics of the Jacobian .
def analyze_jacobian(jacobian_array):
  """Computes eigenvalue statistics of the Jacobian.

  Computes the eigenvalues and condition number of the metric tensor for the
  Jacobian evaluated at each element of the batch and the mean metric tensor
  across the batch.

  Args:
    jacobian_array: A numpy array holding the Jacobian.

  Returns:
    A dict of spectral statistics with two elements, one containing stats
    for every metric tensor in the batch, another for the mean metric tensor.
  """
  # Shape [batch, x_dim, fx_dim].
  jacobian_transpose = np.transpose(jacobian_array, [0, 2, 1])

  # Shape [batch, x_dim, x_dim].
  metric_tensor = np.matmul(jacobian_transpose, jacobian_array)

  mean_metric_tensor = np.mean(metric_tensor, 0)
  # Reshapes to have a dummy batch dimension.
  mean_metric_tensor = np.reshape(mean_metric_tensor,
                                  (1,) + metric_tensor.shape[1:])
  return {
      "metric_tensor": _analyze_metric_tensor(metric_tensor),
      "mean_metric_tensor": _analyze_metric_tensor(mean_metric_tensor)
  }
def analyze_jacobian ( jacobian_array ) : # Shape [batch, x_dim, fx_dim]. jacobian_transpose = np . transpose ( jacobian_array , [ 0 , 2 , 1 ] ) # Shape [batch, x_dim, x_dim]. metric_tensor = np . matmul ( jacobian_transpose , jacobian_array ) mean_metric_tensor = np . mean ( metric_tensor , 0 ) # Reshapes to have a dummy batch dimension. mean_metric_tensor = np . reshape ( mean_metric_tensor , ( 1 , ) + metric_tensor . shape [ 1 : ] ) return { "metric_tensor" : _analyze_metric_tensor ( metric_tensor ) , "mean_metric_tensor" : _analyze_metric_tensor ( mean_metric_tensor ) }