def METHOD_NAME(shape):
    """Computes the number of input and output units for a weight shape.
    Args:
        shape: Integer shape tuple or TF tensor shape.
    Returns:
        A tuple of scalars (fan_in, fan_out).
    """
    if len(shape) < 1:  # Just to avoid errors for constants.
        fan_in = fan_out = 1
    elif len(shape) == 1:
        fan_in = fan_out = shape[0]
    elif len(shape) == 2:
        # Assuming stacked NN.
        # kernel shape: (num_stack, fan_in)
        fan_in = shape[1]
        fan_out = 1
    else:
        # Assuming stacked NN.
        # kernel shape: (..., fan_in, fan_out)
        fan_in = shape[-2]
        fan_out = shape[-1]
    return fan_in, fan_out