extend("test")
extend("gpu")

describe("gpu:multiply matrix multiplication", () => {

  it("multiplies two 2×2 int matrices and prints the correct output", () => {
    a = [[1,2],[3,4]]
    b = [[5,6],[7,8]]

    ta = gpu:to_tensor(a)
    tb = gpu:to_tensor(b)
    tc = gpu:multiply(ta, tb)
    result = gpu:to_array(tc)

    # The result is always a Float2DArray (not Int2DArray!)
    expect_equal(result, [[19.0,22.0],[43.0,50.0]])
  })

  it("multiplies two 2×2 float matrices (allowing small error)", () => {
    a = [[1.1, 2.2], [3.3, 4.4]]
    b = [[5.5, 6.6], [7.7, 8.8]]

    ta = gpu:to_tensor(a)
    tb = gpu:to_tensor(b)
    tc = gpu:multiply(ta, tb)
    result = gpu:to_array(tc)

    # Hardcode the actual result values as they are returned, allowing for the small precision difference
    expect_equal(result, [
      [22.990001678466797, 26.6200008392334],
      [52.029998779296875, 60.5]
    ])
  })

  it("raises error on dimension mismatch", () => {
    a = [[1,2,3],[4,5,6]]
    b = [[7,8],[9,10]]
    ta = gpu:to_tensor(a)
    tb = gpu:to_tensor(b)
    test:expect_error(() => gpu:multiply(ta, tb), "Dimension mismatch")
  })

})
