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

[YodaEmbedding]

An additional flag is too complex. However, we could instead always return y_hat:

class SimpleVAECompressionModel(CompressionModel):
    def forward (self, x):
        ...
        return {
            "x_hat": x_hat,
            "likelihoods": y_out["likelihoods"],
            # For debugging only:
            "y_hat": y_hat,
        }

In some of my own code, I believe I even do **y_out so that I can access everything.

[picagrad]

If that's easier then the flat then I think that's fine.