def _calculate_bilinear_cost(
    op, coeff, num_alive_inputs, num_alive_outputs, batch_size):
  """Calculates bilinear cost for an op.

  Args:
    op: A tf.Operation.
    coeff: A float coefficient for the bilinear function.
    num_alive_inputs: Scalar Tensor indicating how many input channels are
      considered alive.
    num_alive_outputs: Scalar Tensor indicating how many output channels are
      considered alive.
    batch_size: Integer batch size to calculate cost/loss for.

  Returns:
    Tensor with the cost of the op.
  """
  if op.type == 'DepthwiseConv2dNative':
    # num_alive_inputs may not always equals num_alive_outputs because the
    # input (e.g. the image) may not have a gamma regularizer. In this
    # case the computation is proportional only to num_alive_outputs.
    return batch_size * coeff * num_alive_outputs
  else:
    return batch_size * coeff * num_alive_inputs * num_alive_outputs