Bug fix: MultiOutSizeLinear masks outputs to 0 when dimension of parameter >1 and dim * patch_size is not present in out_feat_size.

[gorold]

This PR fixes a major bug for the MultiOutSizeLinear class.

When dim > 1 in the following code snippet,

uni2ts/src/uni2ts/distribution/_base.py

Lines 74 to 85 in c4d75a3

	self.proj = convert_to_module(
	tree_map(
	lambda dim: (
	proj_layer(in_features, dim * out_features, **kwargs)
	if isinstance(out_features, int)
	else proj_layer(
	in_features, tuple(dim * of for of in out_features), **kwargs
	)
	),
	args_dim,
	)
	)

it causes for feat_size in self.out_features_ls be dim * patch_size, whereas out_feat_size contains the patch sizes. This means that certain patch sizes could be ignored.

uni2ts/src/uni2ts/module/ts_embed.py

Lines 172 to 178 in c4d75a3

	for idx, feat_size in enumerate(self.out_features_ls):
	weight = self.weight[idx] * self.mask[idx]
	bias = self.bias[idx] if self.bias is not None else 0
	out = out + (
	torch.eq(out_feat_size, feat_size).unsqueeze(-1)
	* (einsum(weight, x, "out inp, ... inp -> ... out") + bias)
	)

For the current configuration, this means that outputs with patch sizes 8 and 16 are always masked to 0.

Closes #44

[gorold]

lgtm :)