code: def calculate_bias_shape(
      input_shape: Sequence[int], bias_dims: Sequence[int]
  ):
      """Calculate `bias_shape` based on the `input_shape` and `bias_dims`.
  
      Args:
        input_shape: Shape of the input being passed into the module. The leading
          dimension is the mini-batch size.
        bias_dims: The dimensions that bias should be applied over. The remaining
          dimensions will be broadcast over.
  
      Returns:
        bias_shape: Tuple corresponding to the shape of bias Variable to create.
  
      Raises:
        ValueError: If the user attempts to add bias over the mini-batch dimension,
            e.g. `bias_dims=[0]`.
      """
      input_rank = len(input_shape)
      if bias_dims is None:
          # If None, default is to use all dimensions.
          return input_shape[1:]
  
      elif not bias_dims:
          # If empty list, use a scalar bias.
          return ()
  
      else:
          # Otherwise, calculate bias_shape from bias_dims.
          bias_shape = [1] * input_rank
          # Populate bias dimensions.
          for dim in bias_dims:
              if dim < 0:
                  dim %= input_rank
  
              if dim == 0:
                  raise ValueError(
                      "Cannot apply bias across the minibatch dimension."
                  )
              elif dim >= input_rank:
                  raise ValueError(
                      "Dimension %d (bias_dims=%r) out of range for input of rank %r."
                      % (dim, tuple(bias_dims), input_rank)
                  )
  
              bias_shape[dim] = input_shape[dim]
          # Strip leading unit dimensions.
          start = input_rank
          for dim in range(1, input_rank):
              if bias_shape[dim] != 1:
                  start = dim
                  break
          return tuple(bias_shape[start:])
  
file_path: dm-haiku/haiku/_src/bias.py
imports: from collections.abc import Sequence
  from typing import Optional, Union
  from haiku._src import base
  from haiku._src import initializers
  from haiku._src import module
  from haiku._src import utils
  import jax
  import jax.numpy
label_desc: Use image node_id 1 for calling a global function with example usage: calculate_bias_shape(input_shape, bias_dims) and returns: input_shape, , tuple