Allow returning latent space y or y_hat to allow variable rate modification

[picagrad]

Feature

Please allow user to specify whether y_hat (and possibly also y) be returned from the forward and compress methods of all codec families. i.e.:

class SimpleVAECompressionModel.forward(CompressionModel):

	...

	def forward (self, x, return_y_hat:bool = False):

	...

		if return_y_hat:
			return {"x_hat": x_hat, "likelihoods": y_out["likelihoods"], "y_hat": y_hat}
		else:
			return {"x_hat": x_hat, "likelihoods": y_out["likelihoods"]}

Motivation

1. Improves debugging ability for users, often issues can be directly related to shape/values of y_hat
2. Scaling of y_hat can be used to allow simple variable rate adaptation (as in Tong, Kedeng, et al. "Qvrf: A quantization-error-aware variable rate framework for learned image compression." 2023 IEEE International Conference on Image Processing (ICIP). IEEE, 2023.) and access to y_hat is needed to verify/test/debug any such adaptation (although the adaptation itself can be done through workarounds)

Additional context