def model(config: ml_collections.ConfigDict, **kwargs):
return pixelcnn.PixelCNNPP(depth=config.n_resnet,
features=config.n_feature,
logistic_components=config.n_logistic_mix,
**kwargs)
def test_pcnn_shape(self):
x = random.normal(self.rng, (2, 4, 4, 3))
model = pixelcnn.PixelCNNPP(depth=0, features=2, dropout_p=0)
out, _ = model.init_with_output(self.rng, x)
self.assertEqual(out.shape, (2, 4, 4, 100))
def model(**kwargs):
return pixelcnn.PixelCNNPP(
depth=FLAGS.n_resnet, features=FLAGS.n_feature,
logistic_components=FLAGS.n_logistic_mix, **kwargs)
def model(**kwargs):
return pixelcnn.PixelCNNPP(depth=FLAGS.n_resnet,
features=FLAGS.n_feature,
**kwargs)
