def calculate_branching_model_parameters(encoding_depth,
                                         left_inputs,
                                         left_layers,
                                         left_output_dims,
                                         right_inputs,
                                         right_layers,
                                         right_output_dims,
                                         combiner_functions,
                                         layer_registry,
                                         num_cells,
                                         final_combiner_function,
                                         encoder_depth=None,
                                         enforce_output_size=False,
                                         enforce_fixed_output_sizes=True):
  """Calculates the number of parameters in the given model portion.

  Args:
    encoding_depth: Integer. The depth of the initial input tensor.
    left_inputs: Integer list. The indexes of the hidden layer inputs for the
      left branch.
    left_layers: String list. The names of the left branch layers.
    left_output_dims: Integer list. The output dimensions for each of the left
      branch layers.
    right_inputs: Integer list. The indexes of the hidden layer inputs for the
      right branch.
    right_layers: String list. The names of the right branch layers.
    right_output_dims: Integer list. The output dimensions of each of the right
      branch layers.
    combiner_functions: String list. The functions used to combine the left and
      right branch tensors.
    layer_registry: layers.LayerRegistry. The LayerRegistry that contains the
      layers.TranslationLayers needed to construct the model.
    num_cells: Integer. The number of times the given layers are repeated to
      produce the model.
    final_combiner_function: String. The COMBINER_FUNCTIONS key for the combiner
      used to combine the unused hidden dimensions.
    encoder_depth: Integer. The depth of the final encoder layer.
    enforce_output_size: Boolean. If True, include parameters for the addition
      of a dense layer that projects the final output to the appropriate
      `encoding_depth` if it is not already that size. If False, do not add any
      additional parameters.
    enforce_fixed_output_sizes: Whether or not to automatically resize output
      dimensions to match the input dimension if `should_alter_output_dim()`
      returns True.

  Raises:
    ValueError: When the layer config lists are not of equal length.

  Returns:
    total_parameters: The total number of parameters in the model, accounting
      for repeated cells.
    output_depth: The depth of the cell output tensor.
    hidden_depths: The depths of the hidden layers.
    unused_outputs: List of integer indexes of the hidden layers that are not
      used as input, and therefore are concatenated to produce the cell
      output.
  """
  if not (len(left_inputs) == len(left_layers) == len(left_output_dims) ==
          len(right_inputs) == len(right_layers) == len(right_output_dims) ==
          len(combiner_functions)):
    raise ValueError("Layer configs must be of equal length.")

  total_parameters = 0
  output_depth = encoding_depth
  for _ in range(num_cells):
    hidden_depths = [output_depth]
    unused_outputs = set(range(len(left_inputs) + 1))

    for (left_input, left_layer, left_output_dim, right_input,
         right_layer, right_output_dim, combiner_function) in zip(
             left_inputs, left_layers, left_output_dims, right_inputs,
             right_layers, right_output_dims, combiner_functions):

      assert not (left_layer == DEAD_BRANCH_KEY and
                  right_layer == DEAD_BRANCH_KEY)

      if left_layer == DEAD_BRANCH_KEY:
        left_parameters = 0

      else:
        left_input_dim = hidden_depths[left_input]
        if should_alter_output_dim(left_layer, enforce_fixed_output_sizes,
                                   left_input_dim, left_output_dim):
          left_output_dim = left_input_dim

        left_parameters = layer_registry.get(left_layer).num_params(
            left_input_dim, left_output_dim, encoder_depth=encoder_depth)

      if right_layer == DEAD_BRANCH_KEY:
        right_parameters = 0

      else:
        right_input_dim = hidden_depths[right_input]
        if should_alter_output_dim(right_layer, enforce_fixed_output_sizes,
                                   right_input_dim, right_output_dim):
          right_output_dim = right_input_dim

        right_parameters = layer_registry.get(right_layer).num_params(
            right_input_dim, right_output_dim, encoder_depth=encoder_depth)

      total_parameters += left_parameters + right_parameters

      if left_layer == DEAD_BRANCH_KEY:
        hidden_dim = right_output_dim
      elif right_layer == DEAD_BRANCH_KEY:
        hidden_dim = left_output_dim
      else:
        hidden_dim = COMBINER_FUNCTIONS[combiner_function](
        ).combined_output_dim([left_output_dim, right_output_dim])
      hidden_depths.append(hidden_dim)

      try:
        if left_layer != DEAD_BRANCH_KEY:
          unused_outputs.remove(left_input)
      except KeyError:
        pass
      try:
        if right_layer != DEAD_BRANCH_KEY:
          unused_outputs.remove(right_input)
      except KeyError:
        pass

    # All unused outputs combined_together.
    unused_hidden_depths = [hidden_depths[index] for index in unused_outputs]
    output_depth = COMBINER_FUNCTIONS[final_combiner_function](
    ).combined_output_dim(unused_hidden_depths)

  # Add the resizing layer if needed.
  if output_depth != encoding_depth and enforce_output_size:
    total_parameters += layer_registry.get(
        layers.STANDARD_CONV_1X1_REGISTRY_KEY).num_params(
            output_depth, encoding_depth, encoder_depth=encoder_depth)

  return (total_parameters, output_depth, hidden_depths, unused_outputs)
Calculates the number of parameters in the given model portion.

  Args:
    encoding_depth: Integer. The depth of the initial input tensor.
    left_inputs: Integer list. The indexes of the hidden layer inputs for the
      left branch.
    left_layers: String list. The names of the left branch layers.
    left_output_dims: Integer list. The output dimensions for each of the left
      branch layers.
    right_inputs: Integer list. The indexes of the hidden layer inputs for the
      right branch.
    right_layers: String list. The names of the right branch layers.
    right_output_dims: Integer list. The output dimensions of each of the right
      branch layers.
    combiner_functions: String list. The functions used to combine the left and
      right branch tensors.
    layer_registry: layers.LayerRegistry. The LayerRegistry that contains the
      layers.TranslationLayers needed to construct the model.
    num_cells: Integer. The number of times the given layers are repeated to
      produce the model.
    final_combiner_function: String. The COMBINER_FUNCTIONS key for the combiner
      used to combine the unused hidden dimensions.
    encoder_depth: Integer. The depth of the final encoder layer.
    enforce_output_size: Boolean. If True, include parameters for the addition
      of a dense layer that projects the final output to the appropriate
      `encoding_depth` if it is not already that size. If False, do not add any
      additional parameters.
    enforce_fixed_output_sizes: Whether or not to automatically resize output
      dimensions to match the input dimension if `should_alter_output_dim()`
      returns True.

  Raises:
    ValueError: When the layer config lists are not of equal length.

  Returns:
    total_parameters: The total number of parameters in the model, accounting
      for repeated cells.
    output_depth: The depth of the cell output tensor.
    hidden_depths: The depths of the hidden layers.
    unused_outputs: List of integer indexes of the hidden layers that are not
      used as input, and therefore are concatenated to produce the cell
      output.
  
tensor2tensor/models/neural_architecture_search/nas_model.py