def infer_leading_dims(tensor, dim):
    """Param 'dim': number of non-leading dimensions in tensor.
    Returns:
    shape: tensor shape after leading dims
    T: int --size of first leading dim, if two leading dims, o/w 1.
    B: int --size of first leading dim if one, second leading dim if two, o/w 1.
    has_T: boolean --presense of T dim (2 leading)
    has_B: boolean --presence of B dim (at least 1 leading)
    """
    assert tensor.dim() in (dim, dim + 1, dim + 2)
    shape = tensor.shape[-dim:]
    T = B = 1
    has_T = has_B = False
    if tensor.dim() == dim + 2:
        T, B = tensor.shape[:2]
        has_T = has_B = True  # Might have T=1 or B=1.
    elif tensor.dim() == dim + 1:
        B = tensor.shape[0]
        has_B = True
    return shape, T, B, has_T, has_B