ModularM
Modular•13mo ago•
1 reply
Deftioon

An question about compile-time parameters

Hey all,
So I was trying to make a Matrix struct:
#linalg.mojo
from complex import ComplexSIMD

@value
struct Matrix[rows: Int, cols: Int]:
  var data: ComplexSIMD[DType.float64, rows*cols]

  fn __init__(out self, re: Float64, im: Float64):
    self.data = ComplexSIMD[DType.float64, rows*cols](re, im)

When using these Matrices, its guaranteed that rows*cols is a power of two.

I have another struct that uses a Matrix:
struct Foo[T: Int]:
  var state: Matrix[2**T, 1]

And I fill this state with a series of kronecker products of 2x2 Matrices T times, ensuring that the resulting output is definitely of shape [2**T, 1]. That looks like this:
  var filled_matrix = linalg.Matrix[2, 2](0.0, 0.0)
  for i in range(T-1):
      var new_matrix = linalg.kron(filled_matrix, IDENTITY)
      var filled_matrix = new_matrix
  self.state.replace_matrix(filled_matrix)

Here there is an issue. I cannot put filled_matrix into self.state. 
invalid call to 'replace_matrix': method argument #0 cannot be converted from 'Matrix[2, 2]' to 'Matrix[2.__pow__(T), 2.__pow__(T)]'`.

I figured this is because the compiler cannot figure out the shape of filled_matrix, even though it is guaranteed to be a Matrix[2**T, 1]. However, if I turn from compiler-time parameters to dynamic parameters, I am not sure how to create the ComplexSIMD as a struct property.

Is there a way to use ComplexSIMD while also achieving this behaviour? Or will I need to use UnsafePointer instead.
Was this page helpful?