def GTR(NdotHs, p_roughness, gamma=1.):
    """
    Calculate the GTR microfacet distribution coefficient,given the 
    inner products between the surface normals and the half vectors and the 
    surface roughness.

    Inputs:
        NdotHs          NxLx3 torch.tensor containing the inner products
        p_roughness     Nx1 torch.tensor containing the surface roughnesses
    
    Outputs:
        Ds              NxLx1 torch.tensor containing the microfacet distributions
    """
    cosNH2 = (NdotHs ** 2).clamp_(min=0., max=1.)
    p_roughness2 = p_roughness ** 2
    if gamma == 1.:
        cs = (p_roughness2 - 1) / p_roughness2.log()
        Ds = cs / (1 + (p_roughness2 - 1) * cosNH2 + (cosNH2 == 1).float())
        Ds[cosNH2 == 1.] = (-1 / p_roughness2.log() / p_roughness2).repeat(cosNH2.shape[0],1,1)[cosNH2 == 1.]
    else:
        cs = (gamma - 1) * (p_roughness2 - 1) / (1 - p_roughness2 ** (1 - gamma))
        Ds = cs / ((1 + (p_roughness2 - 1) * cosNH2) ** gamma)
    return Ds