title: Efficient matrix-matrix product
body: <p>I'm trying to compute a matrix-matrix product inside a multilevel Bayes model. The data and model variables are shaped as follows:</p>
  
  <pre><code>X_train: (n_samples, n_predictors)
  alpha: (n_groups, )
  beta: (n_predictors, n_groups)
  </code></pre>
  
  <p><code>X_train</code> is dense. The <code>alpha</code>s are the group intercepts and <code>beta</code>s denote the slopes (again, variable for each group). The samples are split in groups; the array <code>group_index</code> (of size <code>(n_samples,)</code>) indicates the group each sample belongs to. In a nutshell, the linear model is</p>
  
  <pre><code>y[n] = alpha[group_index[n]] + &lt; X_train[n, 1:K], beta[1:K, group_index[n]] &gt; 
  </code></pre>
  
  <p>for all <code>n = 1 ... n_samples</code>, where the <code>&lt; , &gt;</code> denotes an inner product.</p>
  
  <p>Here's how implemented this in Python:</p>
  
  <pre class="lang-py prettyprint-override"><code># an element-wise product between two matrices of size [n_samples, n_predictors]
  y = alpha[group_index] + (X_train * beta[:, group_index].T).sum(axis=-1)
  </code></pre>
  
  <p><strong>Question</strong>: Can this implementation be made more memory / cpu efficient? </p>
  
  <p>Here <code>n_samples</code> can run up to 1 million, while <code>n_predictors</code> and <code>n_groups</code> are around 100 or so. What I'm looking for is a vectorized formulation that (1) requires lower storage and (2) runs about as fast compared to what I came up with above.</p>