Add Nx.covariance

[krstopro]

Added Nx.covariance/2 which estimates the covariance matrix from input tensor. Also edited the test for properties of multivariate normal distribution in nx_test.ex.
Few things I am not sure about:

• Currently the function requires the input tensor to be of rank 2. One could make it work for tensors of rank 1 as well, but then I would prefer the user to use Nx.variance.
• Mean can be passed as an option. This is rather convenient when the mean is known or already computed. This is not the approach taken in the rest of Nx (e.g. functions for variance or standard deviation) or NumPy / JAX.

[polvalente]

We should a least accept tensors of shape {..., m, n} where the first k-2 dimensions are batch dimensions.

For the mean, I think we need to receive it as an argument, so we'd have:

def covariance(tensor), do: covariance(tensor, Nx.mean(tensor, axes: [-2]), [])
def covariance(tensor, opts) when is_list(opts), do: covariance(tensor, Nx.mean(tensor, axes: [-2]), opts)
def covariance(tensor, mean), do: covariance(tensor, mean, [])
def covariance(tensor, mean, opts) do
...
end

Although @josevalim will be able to say for certain if passing the mean as an option is really a problem.

[polvalente]

And you'd do mean = to_tensor(mean) right after tensor = to_tensor(tensor)

[josevalim]

Yes, tensors must be arguments. And I would go with this signature:

def covariance(tensor, opts \\ [])
def covariance(tensor, opts) when is_list(opts), do: covariance(tensor, Nx.mean(tensor, axes: [-2]), opts)
def covariance(tensor, mean), do: covariance(tensor, mean, [])
def covariance(tensor, mean, opts) when is_list(opts) do

:)

[krstopro]

I got used to passing arguments as options to the point I forgot Elixir supports multiple-clause functions. 😅

[krstopro]

We should at least accept tensors of shape {..., m, n} where the first k-2 dimensions are batch dimensions.

@polvalente Just to clarify: if we take a batch of datasets as input (shape {..., n, d}), the output should be a batch of covariances (shape {..., d, d}), right?

[polvalente]

Yes

[polvalente]

We can skip these validations because subtract should already apply them for us.

[krstopro]

One of the reasons why I validated the rank of the mean is because the user can pass a single number and then Nx will broadcast it when subtracting. Do we really want this? If the input tensor is of format n x d, then I would prefer the user to hardcode a d-dimensional mean vector.
Agreed for the dimension.

[msluszniak]

There is a function that calculates this in Scholar https://hexdocs.pm/scholar/Scholar.Covariance.html#covariance_matrix/2. Is there any specific reason to implement this once again in Nx?

[krstopro]

@msluszniak I see. I would honestly prefer it in Nx, as the function is rather fundamental and not only necessary for machine learning algorithms. I also thought of using it for unit testing Nx.Random.multivariate_normal (though this might be imported from Scholar as well).

nx/nx/test/nx/random_test.exs

Lines 331 to 337 in f74fcad

	multivariate_normal_centered = Nx.subtract(multivariate_normal, mean)
	assert_all_close(
	Nx.divide(
	Nx.dot(multivariate_normal_centered, [0], multivariate_normal_centered, [0]),
	1000
	),

[polvalente]

There is a function that calculates this in Scholar https://hexdocs.pm/scholar/Scholar.Covariance.html#covariance_matrix/2. Is there any specific reason to implement this once again in Nx?

@msluszniak I didn't know about that one. I believe we should deprecate that in favor of this one in Nx. We could reuse the implementation from Scholar.

@krstopro I see you added a :ddof option as well as the mean argument. Could you check on the possibility of merging both implementations?

[krstopro]

@polvalente @msluszniak Looking at it, but this one seems more general at the moment (e.g. works with batches and complex input). The main difference is in the way they handle bias. I prefer using :biased option as done in Scholar, but I've been following the rest of Nx and used :ddof.

Both implementations contain a bug though :)
If the last axis of input tensor is named, the function will raise an error (as Nx.dot(x, [0], x, [0]) raises an error due to duplicate name). One way to solve this is to extract the names of the tensor, replace the last name with nil using List.replace_at/2, then rename the tensor. Is there a better way?

[polvalente]

I think we should force the output tensor to have no names. So you can rename all axes of x to nil and it should resolve.

[msluszniak]

@msluszniak I didn't know about that one. I believe we should deprecate that in favor of this one in Nx. We could reuse the implementation from Scholar.

Yes, I think it's a good idea to use the implementation from Nx and deprecate the version from Scholar. But to do this, we also need to implement a correlation matrix (which will be trivial as we get covariance).

[krstopro]

I think we should force the output tensor to have no names. So you can rename all axes of x to nil and it should resolve.

Yeah, I guess I should do that. I still like the names on the batch axes though.

Yes, I think it's a good idea to use the implementation from Nx and deprecate the version from Scholar. But to do this, we also need to implement a correlation matrix (which will be trivial as we get covariance).

Agreed.

[polvalente]

I think we should force the output tensor to have no names. So you can rename all axes of x to nil and it should resolve.

Yeah, I guess I should do that. I still like the names on the batch axes though.

The problem is that as we are reducing on a given axis, the semantics start to get fuzzy on what axes should be kept.
I think the simpler solution is to remove all axes' names.

[krstopro]

I think we should force the output tensor to have no names. So you can rename all axes of x to nil and it should resolve.

Yeah, I guess I should do that. I still like the names on the batch axes though.

The problem is that as we are reducing on a given axis, the semantics start to get fuzzy on what axes should be kept. I think the simpler solution is to remove all axes' names.

Indeed, better remove all the names.

[krstopro]

Alright, changed the description and added unit tests for batch inputs. Might be ready to merge.

[josevalim]

💚 💙 💜 💛 ❤️