Cross correlation implemented in Nx?

I didn’t see cross correlation implemented in Nx. However, I believe cross correlation can be implemented in terms of convolution, which is implemented.

cross_corr(f(t), g(t)) = conv(conjugate(f(-t)), g(t))

(If your tensors are real-valued, you can skip conjugate.)

This is a bit of a drive-by response. I’m leaving the implementation unspecified, though hopefully that can get you started.

Also if I’m incorrect about the cross corr. being equivalent to convolution, someone please correct me!

Also, I believe that the definition of convolution in Nx is such that you don’t need to apply reflection to get the correct orientation of the cross-correlation. Feel free to open an issue so we can discuss this further.

Okay, after finding out, that cross corrrelation does not exists. One of the first thoughts of mine were convolution. But I did not try further because it needs three dimension.
@billylanchantin because you said I could use conv I tried it, with adding the needed dimensions.
I tried three things, computing on paper, on elixir and python.
I have no complex data.
In elixir I tried:

test = Nx.tensor([1, 2, 3, 4, 5])
  |> Nx.reshape({1, 1, 5})
Nx.conv(test, Nx.reverse(test))

Returns 35, which is what also computed by hand.

On the other hand with numpy

test = np.array([1, 2, 3, 4, 5])
np.correlate(test, test)
>>>array([55])
np.convolve(test, test)
>>>array([ 1,  4, 10, 20, 35, 44, 46, 40, 25])

Which is strange. Because it seams that array was has not be reversed, but the convolve is the same as I expected. But in the end. But functions can compute the all offsets with the mode=full option.

I tried to implement this behaviour with:

defmodule CrossCorrelation do

  def cross_corr(t1, t2) do
    Nx.conv(Nx.reverse(t1), t2)
  end

  def full_cross_corr(t1, t2) do

    indices1 = Nx.iota(Nx.shape(t1))[0][0]
    coeff1 = Nx.map(indices1, fn x -> 
      ts1 = t1[0 .. Nx.to_number(x)]  # error here: Nx.Defn.Expr.to_binary/2
      ts2 = t2[0 .. Nx.to_number(x)]
      Nx.conv(Nx.reverse(ts1), ts2)
    end)

    indices2 = Nx.iota(Nx.shape(t1))[0][0][1 .. -1 //1]
    coeff2 = Nx.map(indices2, fn x -> 
      ts1 = t1[x .. -1//1]
      ts2 = t2[x .. -1//1]
      Nx.conv(Nx.reverse(ts1), ts2)
    end)

    Nx.concatenate([coeff1, coeff2])
  end


  def find_offset(t1, t2) do
    l1 = elem(Nx.shape(t1), 0)
    l1 - Nx.argmax(full_cross_corr(t1, t2)) - 1
  end
end

But the I could not test it, as the function full_cross_corr raises breaks with the error:
** (ArgumentError) cannot invoke to_binary/2 on Nx.Defn.Expr. From within the Nx.to_number function.
But not to capsule x with Nx.to_number throws:

** (ArgumentError) ranges (first..last) expect both sides to be integers, got: 0..#Nx.Tensor<
  s64
  
  Nx.Defn.Expr
  parameter a:0   s64
>

So x is tensor of type integer with value 0. But what is the problem with Nx.to_number than?

This implementation isn’t advisable. Please open an issue so we can discuss implementing cross_correlation in Nx itself.

general problems:

• Nx.map receives a function that operates on Nx.Defn.Expr (basically, it will build a computation), so it doesn’t have acess to the tensor values – this is the problem with Nx.to_number
• Nx.map isn’t really optimizeable by compilers and should ideally be avoided. If cross correlation can’t be implemented with the current tooling, we should add a new function for supporting it.
• You should focus on implementing functions with Nx using defn as widely as possible, because defn will generate computation graphs that can be then compiled to native machine code in a more efficient way

In your issue, please describe in general what you expect that the cross correlation function should do, with example inputs and their corresponding outputs.

If there’s similar functionality in Jax or PyTorch, a link would be helpful!
I’m generally familiar with cross-correlation, convolution and so on, but there are many conflicting definitions for those mathematical operations.

I think, I can now better articulate a problem I have with Nx.conv.

# Define the two time series
t1 = np.array([0.1, 0.2, 0.4, 0.7, 0.9, 0.8, 0.6, 0.3, 0.1])
t2 = np.array([0.3, 0.6, 0.8, 0.7, 0.4, 0.2, 0.1, 0.1, 0.3])

In Nx I added |> Nx.reshape({1, 1, 9})
I just compared the results from Numpy and Nx:

Should is that correct? Should I open a ticket?

I updated my implemention. Because I had some problem with Nx. I used more standard Elixir.

defmodule CrossCorrelation do
  def cross_corr(t1, t2) do
    Nx.conv(t1, t2)
  end

  def full_cross_corr(t1, t2) do
    l1 = elem(Nx.shape(t1), 2) - 1
    indices1 = 0 .. l1

    coeff1 =
      Enum.map(indices1, fn x ->

        ts1 = t1[[.., .., 0..x]]
        ts2 = t2[[.., .., 0..x]]
        Nx.to_number(Nx.conv(ts1, ts2)[0][0][0])
      end)

    indices2 = 1..l1

    coeff2 =
      Enum.map(indices2, fn x ->
        ts1 = t1[[.., .., x..-1//1]]
        ts2 = t2[[.., .., x..-1//1]]
        Nx.to_number(Nx.conv(ts1, ts2)[0][0][0])
      end)
   
    coeff1 ++ coeff2
  end

  def find_offset(t1, t2) do
    l1 = elem(Nx.shape(t1), 2)

    coeffs = full_cross_corr(t1, t2)
    Nx.argmax(Nx.tensor(coeffs))
    Nx.to_number(Nx.argmax(Nx.tensor(coeffs))) - l1 + 1
  end
end

Problem I have is: that the result is good when compare the simple autocorrelation with the Tensor [[[1, 2, 3, 4, 5]]]. CrossCorrelation.full_cross_corr compares good to numpy.correlate(t1, t2, mode=full). But for the more realistic, time series t1 and t2. The result differs greatly.
I do not know, whether I did something wrong, or its something different. There is not something similar to numpy.correlate(t1, t2, mode=full) yet in Nx, right?

P.S.: We wrote at the same time. I added some links the table.
I also wanted to add some test results:

numpy:

cross_corr = np.correlate(t1, t2, mode='full')

Returns:
0.03,0.07,0.15,0.29000000000000004,0.46,0.6300000000000001,0.8699999999999999,1.1900000000000002,1.6,1.9699999999999998,2.1300000000000003,1.9500000000000002,1.48,0.9099999999999999,0.44,0.15,0.03

My Elixir implementation:

def full_cross_corr(t1, t2) do
    l1 = elem(Nx.shape(t1), 2) - 1
    indices1 = 0 .. l1

    coeff1 =
      Enum.map(indices1, fn x ->

        ts1 = t1[[.., .., 0..x]]
        ts2 = t2[[.., .., 0..x]]
        Nx.to_number(Nx.conv(ts1, ts2)[0][0][0])
      end)

    indices2 = 1..l1

    coeff2 =
      Enum.map(indices2, fn x ->
        ts1 = t1[[.., .., x..-1//1]]
        ts2 = t2[[.., .., x..-1//1]]
        Nx.to_number(Nx.conv(ts1, ts2)[0][0][0])
      end)
   
    coeff1 ++ coeff2
  end

returns:
[0.030000001192092896, 0.15000000596046448, 0.4700000286102295, 0.9600000381469727,
1.3200000524520874, 1.4800000190734863, 1.5399999618530273, 1.5699999332427979, 1.600000023841858,
1.5699999332427979, 1.4500000476837158, 1.1299999952316284, 0.6399999856948853, 0.2800000309944153,
0.12000000476837158, 0.06000000238418579, 0.030000001192092896]

Yeah, so this is one of the issues I meant with conflicting definitions.
Mathematically, convolution and correlation are the same idea, a series of sliding inner products, with the difference being that in one of those you reflect one of the tensors and in the other, you don’t.

The thing is, in computational graphics, the convolution operation tends to not reflect the kernel, so it really is a correlation. This is why you’re getting “swapped” results there.

I believe if you use Nx.reverse on the second operand (the kernel) and use padding: :same for the convolution, you’ll get the results you’re looking for.

Ok, so after some analysis, the call you’re looking for needs to add some zero-padding to your tensors:

iex(57)> Nx.conv(Nx.reshape(t1, {1, 1, 9}), Nx.reshape(t2, {1, 1, 9}), padding: [{8, 8}])
#Nx.Tensor<
  f32[1][1][17]
  Torchx.Backend(cpu)
  [
    [
      [0.030000001192092896, 0.07000000029802322, 0.15000000596046448, 0.2900000214576721, 0.46000003814697266, 0.6299999952316284, 0.8700000643730164, 1.190000057220459, 1.6000001430511475, 1.9700000286102295, 2.130000114440918, 1.9500000476837158, 1.4800000190734863, 0.9100000262260437, 0.4400000274181366, 0.15000000596046448, 0.030000001192092896]
    ]
  ]
>

If you look closely, the result from convolving with Nx.conv(..., padding: :same) is contained in the middle of your results from numpy

Yes, it is! I also saw, that same padding: returns the inner part of what I wanted.

1. Is there a shorter version for adding the extra dim to the tensors t1 and t2.
2. How would I select the correct padding? I assume its innerTensorDim - 1.
3. Could you add an atom for this padding to Nx.conv? But Tensorflow also only has valid and same. While Numpy-version I showed has: full, same, valid. But the are made for different things. I also just “missused” Nx.conv ^^

Update:
@polvalente which your help I could solve the problem. Thanks alot!

I think the padding used can be size t2 - 1, not sure about that. You’d have to check what happens when the tensors are of different sizes. It might actually work to do size1 + size2 - 1 on a single side only.

As for adding the axes, you can use Nx.new_axis