Defining Nx tensors at compile time - recommendations requested

In a new color library I’m writing there is a lot of matrix math to do with chromaticity adaptation, conversion between linear space and gamma space and conversions between color spaces in general.

Many of these conversion matrices can be computed at compile time and Nx is a great fit for this and is working well. Here’s an example:

  @chromatic_adaptations %{
    xyz_scaling: %{
      matrix: [
        [1.0000000,  0.0000000,  0.0000000],
        [0.0000000,  1.0000000,  0.0000000],
        [0.0000000,  0.0000000,  1.0000000]
      ] |> Nx.tensor(),

      inverse_matrix: [
        [1.0000000,  0.0000000,  0.0000000],
        [0.0000000,  1.0000000,  0.0000000],
        [0.0000000,  0.0000000,  1.0000000]
      ] |> Nx.tensor()
    },
    ...
}

However by doing this the tensors are not “backend” aware. Color shouldn’t impose any configuration on the user at all. Therefore I’m not sure what the right strategy is to define tensors at compile time and yet still fit into a users Nx backend configuration.

One possible approach might be to:

  • Pre-compute the tensors
  • Serialize the data
  • Load the data at application start and Nx.backend_transfer/2 the data
  • Store the resulting tensors in :persistent_term

Thoughts and suggestions very welcome.

4 Likes

Looking at the code I think what you want to do is turn most or all of the functions in chromatic adaptations into defn

Then I would explicitly set the backend of the compile time constants to Nx.BinaryBackend (this will prevent compile issues if you have set a different backend already)

Implementing functions as defn will make the implementations backend aware by default

3 Likes

Thanks @seanmor5, sounds like a plan. Much appreciated.

1 Like