`ch5/1.open-information-extraction.ipynb` - Listing 5.1 `RuntimeError: could not create a primitive descriptor for a matmul primitive` on Apple M2

[tmancill]

Hi, when trying to run Listing 5.1 via the provided docker-compose image on an Apple M2 machine, I get an exception (below). Other sections of the notebooks are working as expected, and the same code works correctly on an x86_64-based Linux system, also running via docker-compose. Perhaps an issue with spaCy or PyTorch support on that architecture? (I can offer to help test patches, but I'm not sure how to further triage the underlying issue.)

Specific notebook entry:

https://github.com/treygrainger/ai-powered-search/blob/4e1681db4992cc4cafb10d3780723f9c32109096/docker/data-science/notebooks/ch05/1.open-information-extraction.ipynb#L167-L186

exception

---------------------------------------------------------------------------
RuntimeError                              Traceback (most recent call last)
Cell In[8], line 16
     13         facts.extend(resolve_facts(sentence)) # subj:(Companies), rel:(employ), obj:(Data Scientists)
     14     return facts
---> 16 graph = generate_graph(text)
     17 for i in graph: print(i)

Cell In[8], line 7, in generate_graph(text)
      6 def generate_graph(text):
----> 7     doc = coref_model(text)    
      8     doc = resolve_coreferences(doc) # "they" => "Data Scientists"
      9     sentences = get_sentences(lang_model(doc)) # Data Scientists also write code. => ['nsubj, 'advmod', ROOT', 'dobj', 'punct']    

File /opt/conda/lib/python3.10/site-packages/spacy/language.py:1031, in Language.__call__(self, text, disable, component_cfg)
   1029     raise ValueError(Errors.E109.format(name=name)) from e
   1030 except Exception as e:
-> 1031     error_handler(name, proc, [doc], e)
   1032 if not isinstance(doc, Doc):
   1033     raise ValueError(Errors.E005.format(name=name, returned_type=type(doc)))

File /opt/conda/lib/python3.10/site-packages/spacy/util.py:1670, in raise_error(proc_name, proc, docs, e)
   1669 def raise_error(proc_name, proc, docs, e):
-> 1670     raise e

File /opt/conda/lib/python3.10/site-packages/spacy/language.py:1026, in Language.__call__(self, text, disable, component_cfg)
   1024     error_handler = proc.get_error_handler()
   1025 try:
-> 1026     doc = proc(doc, **component_cfg.get(name, {}))  # type: ignore[call-arg]
   1027 except KeyError as e:
   1028     # This typically happens if a component is not initialized
   1029     raise ValueError(Errors.E109.format(name=name)) from e

File /opt/conda/lib/python3.10/site-packages/spacy/pipeline/trainable_pipe.pyx:56, in spacy.pipeline.trainable_pipe.TrainablePipe.__call__()

File /opt/conda/lib/python3.10/site-packages/spacy/util.py:1670, in raise_error(proc_name, proc, docs, e)
   1669 def raise_error(proc_name, proc, docs, e):
-> 1670     raise e

File /opt/conda/lib/python3.10/site-packages/spacy/pipeline/trainable_pipe.pyx:52, in spacy.pipeline.trainable_pipe.TrainablePipe.__call__()

File /opt/conda/lib/python3.10/site-packages/spacy_experimental/coref/coref_component.py:153, in CoreferenceResolver.predict(self, docs)
    150     out.append([])
    151     continue
--> 153 scores, idxs = self.model.predict([doc])
    154 # idxs is a list of mentions (start / end idxs)
    155 # each item in scores includes scores and a mapping from scores to mentions
    156 ant_idxs = idxs

File /opt/conda/lib/python3.10/site-packages/thinc/model.py:334, in Model.predict(self, X)
    330 def predict(self, X: InT) -> OutT:
    331     """Call the model's `forward` function with `is_train=False`, and return
    332     only the output, instead of the `(output, callback)` tuple.
    333     """
--> 334     return self._func(self, X, is_train=False)[0]

File /opt/conda/lib/python3.10/site-packages/thinc/layers/chain.py:54, in forward(model, X, is_train)
     52 callbacks = []
     53 for layer in model.layers:
---> 54     Y, inc_layer_grad = layer(X, is_train=is_train)
     55     callbacks.append(inc_layer_grad)
     56     X = Y

File /opt/conda/lib/python3.10/site-packages/thinc/model.py:310, in Model.__call__(self, X, is_train)
    307 def __call__(self, X: InT, is_train: bool) -> Tuple[OutT, Callable]:
    308     """Call the model's `forward` function, returning the output and a
    309     callback to compute the gradients via backpropagation."""
--> 310     return self._func(self, X, is_train=is_train)

File /opt/conda/lib/python3.10/site-packages/spacy_experimental/coref/coref_model.py:85, in coref_forward(model, X, is_train)
     84 def coref_forward(model: Model, X, is_train: bool):
---> 85     return model.layers[0](X, is_train)

File /opt/conda/lib/python3.10/site-packages/thinc/model.py:310, in Model.__call__(self, X, is_train)
    307 def __call__(self, X: InT, is_train: bool) -> Tuple[OutT, Callable]:
    308     """Call the model's `forward` function, returning the output and a
    309     callback to compute the gradients via backpropagation."""
--> 310     return self._func(self, X, is_train=is_train)

File /opt/conda/lib/python3.10/site-packages/thinc/layers/pytorchwrapper.py:225, in forward(model, X, is_train)
    222 convert_outputs = model.attrs["convert_outputs"]
    224 Xtorch, get_dX = convert_inputs(model, X, is_train)
--> 225 Ytorch, torch_backprop = model.shims[0](Xtorch, is_train)
    226 Y, get_dYtorch = convert_outputs(model, (X, Ytorch), is_train)
    228 def backprop(dY: Any) -> Any:

File /opt/conda/lib/python3.10/site-packages/thinc/shims/pytorch.py:97, in PyTorchShim.__call__(self, inputs, is_train)
     95     return self.begin_update(inputs)
     96 else:
---> 97     return self.predict(inputs), lambda a: ...

File /opt/conda/lib/python3.10/site-packages/thinc/shims/pytorch.py:115, in PyTorchShim.predict(self, inputs)
    113 with torch.no_grad():
    114     with torch.cuda.amp.autocast(self._mixed_precision):
--> 115         outputs = self._model(*inputs.args, **inputs.kwargs)
    116 self._model.train()
    117 return outputs

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /opt/conda/lib/python3.10/site-packages/spacy_experimental/coref/pytorch_coref_model.py:88, in CorefClusterer.forward(self, word_features)
     85     top_rough_scores_batch = top_rough_scores[i : i + batch_size]
     87     # a_scores_batch    [batch_size, n_ants]
---> 88     a_scores_batch = self.ana_scorer(
     89         all_mentions=words,
     90         mentions_batch=words_batch,
     91         pairwise_batch=pairwise_batch,
     92         top_indices_batch=top_indices_batch,
     93         top_rough_scores_batch=top_rough_scores_batch,
     94     )
     95     a_scores_lst.append(a_scores_batch)
     97 coref_scores = torch.cat(a_scores_lst, dim=0)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /opt/conda/lib/python3.10/site-packages/spacy_experimental/coref/pytorch_coref_model.py:165, in AnaphoricityScorer.forward(self, all_mentions, mentions_batch, pairwise_batch, top_indices_batch, top_rough_scores_batch)
    160 pair_matrix = self._get_pair_matrix(
    161     all_mentions, mentions_batch, pairwise_batch, top_indices_batch
    162 )
    164 # [batch_size, n_ants]
--> 165 scores = top_rough_scores_batch + self._ffnn(pair_matrix)
    166 scores = add_dummy(scores, eps=True)
    168 return scores

File /opt/conda/lib/python3.10/site-packages/spacy_experimental/coref/pytorch_coref_model.py:175, in AnaphoricityScorer._ffnn(self, x)
    170 def _ffnn(self, x: torch.Tensor) -> torch.Tensor:
    171     """
    172     x: tensor of shape (batch_size x rough_k x n_features
    173     returns: tensor of shape (batch_size x antecedent_limit)
    174     """
--> 175     x = self.out(self.hidden(x))
    176     return x.squeeze(2)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1518, in Module._wrapped_call_impl(self, *args, **kwargs)
   1516     return self._compiled_call_impl(*args, **kwargs)  # type: ignore[misc]
   1517 else:
-> 1518     return self._call_impl(*args, **kwargs)

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/module.py:1527, in Module._call_impl(self, *args, **kwargs)
   1522 # If we don't have any hooks, we want to skip the rest of the logic in
   1523 # this function, and just call forward.
   1524 if not (self._backward_hooks or self._backward_pre_hooks or self._forward_hooks or self._forward_pre_hooks
   1525         or _global_backward_pre_hooks or _global_backward_hooks
   1526         or _global_forward_hooks or _global_forward_pre_hooks):
-> 1527     return forward_call(*args, **kwargs)
   1529 try:
   1530     result = None

File /opt/conda/lib/python3.10/site-packages/torch/nn/modules/linear.py:114, in Linear.forward(self, input)
    113 def forward(self, input: Tensor) -> Tensor:
--> 114     return F.linear(input, self.weight, self.bias)

RuntimeError: could not create a primitive descriptor for a matmul primitive