Redundancy of jnp.where operation in d3pm p_logits.

[cloneofsimo]

google-research/d3pm/images/diffusion_categorical.py

Lines 448 to 452 in 52222c0

	model_logits = jnp.where(t_broadcast == 0,
	pred_x_start_logits,
	self.q_posterior_logits(pred_x_start_logits, x,
	t, x_start_logits=True)
	)

Here, we can see that p_logit function is returning logit of prediction of x_0 if t == 0, otherwise q_posterior. which is understandable,

However, looking at the q_posterior_logits method, we can see that it is already returning x_start if t is 0.

google-research/d3pm/images/diffusion_categorical.py

Lines 423 to 425 in 52222c0

	return jnp.where(t_broadcast == 0, tzero_logits,
	out)

So I think we don't need this LOC, am I missing something?