<commit_before><commit_msg>Add matrix_multiply test.<commit_after>#include "array.h"
#include "test.h"

namespace array {

typedef shape<dim<>, dense_dim<>> matrix_shape;

template <typename T>
using matrix = array<T, shape<dim<>, dense_dim<>>>;

template <typename T>
matrix<T> make_matrix(int M, int N) {
  matrix_shape shape({0, M, N}, {0, N});
  return matrix<T>(shape, 0);
}

template <typename T>
index_t rows(const matrix<T>& m) { return m.shape().template dim<0>().extent(); }
template <typename T>
index_t cols(const matrix<T>& m) { return m.shape().template dim<1>().extent(); }

template <typename T>
void multiply_naive(const matrix<T>& a, const matrix<T>& b, matrix<T>& c) {
  for (int i = 0; i < rows(c); i++) {
    for (int j = 0; j < cols(c); j++) {
      c(i, j) = 0;
      for (int k = 0; k < cols(a); k++) {
        c(i, j) += a(i, k) * b(k, j);
      }
    }
  }
}

template <typename T>
void multiply(const matrix<T>& a, const matrix<T>& b, matrix<T>& c) {
  for (int i = 0; i < rows(c); i++) {
    for (int j = 0; j < cols(c); j++) {
      c(i, j) = 0;
    }
    for (int k = 0; k < cols(a); k++) {
      for (int j = 0; j < cols(c); j++) {
        c(i, j) += a(i, k) * b(k, j);
      }
    }
  }
}

TEST(matrix_multiply) {
  auto a = make_matrix<float>(10, 20);
  auto b = make_matrix<float>(20, 30);
  auto c = make_matrix<float>(10, 30);
  a.for_each_value([](float& x) { x = rand(); });
  b.for_each_value([](float& x) { x = rand(); });
  multiply(a, b, c);

  auto c_ref = make_matrix<float>(10, 30);
  multiply_naive(a, b, c_ref);
  
  for (int i = 0; i < 10; i++) {
    for (int j = 0; j < 30; j++) {
      ASSERT_EQ(c(i, j), c_ref(i, j));
    }
  }
}

}  // namespace array
{'fraction_non_alphanumeric': 0.2143276570757487, 'fraction_numerical': 0.021139166177334117, 'mean_word_length': 3.0961538461538463, 'pattern_counts': {'":': 0, '<': 41, '<?xml version=': 0, '>': 32, 'https://': 0, 'lorem ipsum': 0, 'www.': 0, 'xml': 0}, 'pii_count': 0, 'substrings_counts': 0, 'word_list_counts': {'cursed_substrings.json': 0, 'profanity_word_list.json': 0, 'sexual_word_list.json': 0, 'zh_pornsignals.json': 0}}
{'dir': 'git-commits-cleaned', 'id': '4478770', 'n_tokens_mistral': 742, 'n_tokens_neox': 708, 'n_words': 199}