Expand the tokenization tree writing to support per transform block
type coding in next CLs.

Change-Id: I3560e658f89cc500eb49603f95dd2b4e99045f5b