Nx With Custom Number System

I wrote a library for galois field arithmetic and would like to use that in combination with Nx tensors.
All arithmetic operations (+,-,*,/) would have to happen within the field, while transformations like transpose do not.
So for example:

# field with 2^2 values (0,1,2,3)
gf = GF(2,2)
t = Nx.tensor([2,2], backend: gf)
# expected: #Nx.Tensor<GF[2] [2, 2]>

Nx.add(t, t) 
# expected: #Nx.Tensor<GF[2] [0, 0]>
# --> (2 + 2) mod 4 = 0 

Can I use Nx.Backend for that? It seems like a lot of work and I am hoping there is a simpler way :sweat_smile:

1 Like

You can implement the few operations which have unique definitions as a separate backend and then defer to the default behavior (BinaryBackend) or raise on unsupported functions for all of the others!

1 Like

Alternatively, you can define a library which implements the operations on top of existing Nx ops as defn and then the ops are automatically compatible with accelerated backends and compilers

1 Like

Thanks, I will look into both options and may come around with more questions

1 Like

I was experimenting with the second option so far and have a follow up question.
In prime extension fields numbers are represented as polynomials and all calculations are done using these polynomials.
In GF(2^2):

0 = 0     ->  (0*x + 0)  -> [0, 0]
1 = 1     ->  (0*x + 1)  -> [0, 1]
2 = x     ->  (1*x + 0)  -> [1, 0]
3 = x + 1 ->  (1*x + 1)  -> [1, 1]

Therefore 2 * 2 is

2 * 2
= (1*x + 0) * (1*x + 0)   -> [1, 0] * [1, 0]
= x * x 
= x^2                     -> [1, 0, 0]

How could I achieve this behaviour with Nx functions?


There is element-wise multiplication, but I need that “overflow” into the next degree.
I found the operation Nx. outer(t1, t2), which can give me a matrix of element * row results like so:

# [1,0] * [1,0]
Nx.outer(Nx.tensor([1,0]), Nx.tensor([1,0]))
#Nx.Tensor<
  s64[2][2]
  [
    [1, 0], # 1 * [1,0]
    [0, 0]  # 0 * [1,0]
  ]
>

If I can pad and “shift” the first row to the left I can sum the tensor over the y axis:

# 2 * 2   =   x * x   =   x^2   =  [1,0,0]
Nx.outer(Nx.tensor([1,0]), Nx.tensor([1,0]))
    # [1, 0],
    # [0, 0]
|> Nx.pad(0, [{0,0,0}, {1,0,0}]) # insert single zero into all rows
    # [0, 1, 0],
    # [0, 0, 0]
|> shift([1,0]) # <-- made up: cycle first row by one to the left, second not at all
    # [1, 0, 0],
    # [0, 0, 0]
|> Nx.sum(axis: 1)
    # [1, 0, 0]

another example:

# 3 * 2   =   (x + 1) * x   =   x^2 + x  =  [1,1,0]
Nx.outer(Nx.tensor([1,1]), Nx.tensor([1,0]))
    # [1, 0], 
    # [1, 0]
|> Nx.pad(0, [{0,0,0}, {1,0,0}]) # insert single zero into all rows
    # [0, 1, 0],
    # [0, 1, 0]
|> shift([1,0]) # <-- made up: cycle first row by one to the left, second not at all
    # [1, 0, 0],
    # [0, 1, 0]
|> Nx.sum(axis: 1)
    # [1, 1, 0]

Is there a function that behaves like shift?

It’s been a while since I looked at Nx (funny to be saying this) but I think in Nx you could use protocols to implement your own operations over your own types… I would think that if you’re doing GF math you might want to implement the numbers as custom structs instead of trying to use tensors of Nx integers with a different backend?

I found protocols for Container and Stream which are data-structure related protocols. Container may be interesting further down the line though!

In the end all numbers in GF are polynomials, so storing the numbers as tensors of coefficients seems like the best approach to me. Some calculations will be required with these polynomials, so converting from and to a GF digit every time may have a performance impact as well.
The only downside I found so far is the wasted space inside a tensor, storing numbers from 0…3 in a u8 for example.

Unless you mean custom structs as in custom number types (u4 or so), that would be awesome, not just for this topic :smiley: