[Haskell-cafe] type level function with different return kinds

[Zoran Bošnjak]

Dear haskellers,
is it possible to write a type/kind level function that returns types of 
different kinds for different inputs. Either by using open/closed type 
families, a typeclass, extensions or any other type level magic. 
Something like this:

{-# LANGUAGE DataKinds #-}
{-# LANGUAGE TypeFamilies #-}

import GHC.TypeLits

data A = A1 | A2
data B = B1 | B2

type family F (n :: Nat) where
     F 0 = 'A1
     F 1 = 'A2
     F 2 = 'B1
     F 3 = 'B2

The problem with this approach is that (F 0) and (F 1) are ok, but not 
in combination with (F 2). The error says:
     • Expected kind ‘A’, but ‘'B1’ has kind ‘B’

If I instead write:
type family F (n :: Nat) :: k where
... I get compile error already on the (F 0) line
     • Expected kind ‘k’, but ‘'A1’ has kind ‘A’

Any ideas how to remove the "same kind" restriction or how to workaround 
it?

I am using ghc 9.6.6.

regards,
Zoran