def METHOD_NAME(
    network: snt.Module,
    input_spec: List[Union[types.NestedSpec, tf.TensorSpec]],
) -> Optional[tf.TensorSpec]:
  """Builds the network with dummy inputs to create the necessary variables.
  Args:
    network: Sonnet Module whose variables are to be created.
    input_spec: list of input specs to the network. The length of this list
      should match the number of arguments expected by `network`.
  Returns:
    output_spec: only returns an output spec if the output is a tf.Tensor, else
        it doesn't return anything (None); e.g. if the output is a
        tfp.distributions.Distribution.
  """
  # Create a dummy observation with no batch dimension.
  dummy_input = zeros_like(input_spec)
  # If we have an RNNCore the hidden state will be an additional input.
  if isinstance(network, snt.RNNCore):
    initial_state = squeeze_batch_dim(network.initial_state(1))
    dummy_input += [initial_state]
  # Forward pass of the network which will create variables as a side effect.
  dummy_output = network(*add_batch_dim(dummy_input))
  # Evaluate the input signature by converting the dummy input into a
  # TensorSpec. We then save the signature as a property of the network. This is
  # done so that we can later use it when creating snapshots. We do this here
  # because the snapshot code may not have access to the precise form of the
  # inputs.
  input_signature = tree.map_structure(
      lambda t: tf.TensorSpec((None,) + t.shape, t.dtype), dummy_input)
  network._input_signature = input_signature  # pylint: disable=protected-access
  def spec(output):
    # If the output is not a Tensor, return None as spec is ill-defined.
    if not isinstance(output, tf.Tensor):
      return None
    # If this is not a scalar Tensor, make sure to squeeze out the batch dim.
    if tf.rank(output) > 0:
      output = squeeze_batch_dim(output)
    return tf.TensorSpec(output.shape, output.dtype)
  return tree.map_structure(spec, dummy_output)