code: def split_last_dimension(x, n):
    
    x_shape = common_layers.shape_list(x)
    m = x_shape[-1]
    if isinstance(m, int) and isinstance(n, int):
      assert m % n == 0
    return tf.reshape(x, x_shape[:-1] + [n, m // n])
docs: Reshape x so that the last dimension becomes two dimensions.
  
    The first of these two dimensions is n.
  
    Args:
      x: a Tensor with shape [..., m]
      n: an integer.
  
    Returns:
      a Tensor with shape [..., n, m/n]