code: def _compute_fans(shape, fan_in_axes=None):
      """Computes the number of input and output units for a weight shape."""
      if len(shape) < 1:
          fan_in = fan_out = 1
      elif len(shape) == 1:
          fan_in = fan_out = shape[0]
      elif len(shape) == 2:
          fan_in, fan_out = shape
      else:
          if fan_in_axes is not None:
              # Compute fan-in using user-specified fan-in axes.
              fan_in = np.prod([shape[i] for i in fan_in_axes])
              fan_out = np.prod(
                  [
                      s
                      for i, s in enumerate(shape)
                      if i not in fan_in_axes
                  ]
              )
          else:
              # If no axes specified, assume convolution kernels (2D, 3D, or more.)
              # kernel_shape: (..., input_depth, depth)
              receptive_field_size = np.prod(shape[:-2])
              fan_in = shape[-2] * receptive_field_size
              fan_out = shape[-1] * receptive_field_size
      return fan_in, fan_out
  
imports: from collections.abc import Sequence
  from typing import Any, Union
  from haiku._src import base
  from haiku._src.typing import Initializer
  import jax
  import jax.numpy
  import numpy
label_desc: Use image node_id 1 for calling a global function with example usage: _compute_fans(shape, fan_in_axes) and returns: fan_in,  fan_out