pymc-devs / pymc-devs/pytensor

Batched dot doesn't support complex inputs

Open
#1,849 1 comment 0 reactions 0 assignees View on GitHub

Nobody has claimed this yet.

help wanted linalg NumPy compatibility
Dominant language
Python
Stars
644
Forks
208
Avg merge
2d 14h
Merged PRs (30d)
16

Description

Description

This came up in #1840.

import pytensor
import pytensor.tensor as pt
import numpy as np

x = pt.ztensor3('x')
y = pt.ztensor3('y')
z = x @ y
fn = pytensor.function([x, y], z)

x_val, y_val = np.random.normal(size=(2, 6, 6, 2)).view(np.complex128)
fn(x_val, y_val)
Traceback
---------------------------------------------------------------------------
NotImplementedError                       Traceback (most recent call last)
File ~/Documents/Python/pytensor/pytensor/compile/function/types.py:1038, in Function.__call__(self, output_subset, *args, **kwargs)
   1037 try:
-> 1038     outputs = vm() if output_subset is None else vm(output_subset=output_subset)
   1039 except Exception:

NotImplementedError: type(x) is not double or float

During handling of the above exception, another exception occurred:

NotImplementedError                       Traceback (most recent call last)
Cell In[10], line 11
      8 fn = pytensor.function([x, y], z)
     10 x_val, y_val = np.random.normal(size=(2, 6, 6, 2)).view(np.complex128)
---> 11 fn(x_val, y_val)

File ~/Documents/Python/pytensor/pytensor/compile/function/types.py:1048, in Function.__call__(self, output_subset, *args, **kwargs)
   1046     if hasattr(self.vm, "thunks"):
   1047         thunk = self.vm.thunks[self.vm.position_of_error]
-> 1048     raise_with_op(
   1049         self.maker.fgraph,
   1050         node=self.vm.nodes[self.vm.position_of_error],
   1051         thunk=thunk,
   1052         storage_map=getattr(self.vm, "storage_map", None),
   1053     )
   1054 else:
   1055     # old-style linkers raise their own exceptions
   1056     raise

File ~/Documents/Python/pytensor/pytensor/link/utils.py:526, in raise_with_op(fgraph, node, thunk, exc_info, storage_map)
    521     warnings.warn(
    522         f"{exc_type} error does not allow us to add an extra error message"
    523     )
    524     # Some exception need extra parameter in inputs. So forget the
    525     # extra long error message in that case.
--> 526 raise exc_value.with_traceback(exc_trace)

File ~/Documents/Python/pytensor/pytensor/compile/function/types.py:1038, in Function.__call__(self, output_subset, *args, **kwargs)
   1036     t0_fn = time.perf_counter()
   1037 try:
-> 1038     outputs = vm() if output_subset is None else vm(output_subset=output_subset)
   1039 except Exception:
   1040     self._restore_defaults()

NotImplementedError: type(x) is not double or float
Apply node that caused the error: BatchedDot(x, y)
Toposort index: 0
Inputs types: [TensorType(complex128, shape=(None, None, None)), TensorType(complex128, shape=(None, None, None))]
Inputs shapes: [(6, 6, 1), (6, 6, 1)]
Inputs strides: [(96, 16, 16), (96, 16, 16)]
Inputs values: ['not shown', 'not shown']
Outputs clients: [[output[0](BatchedDot.0)]]

Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
  File "/Users/jessegrabowski/mambaforge/envs/pytensor-dev/lib/python3.12/site-packages/ipykernel/zmqshell.py", line 549, in run_cell
    return super().run_cell(*args, **kwargs)
  File "/Users/jessegrabowski/mambaforge/envs/pytensor-dev/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3098, in run_cell
    result = self._run_cell(
  File "/Users/jessegrabowski/mambaforge/envs/pytensor-dev/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3153, in _run_cell
    result = runner(coro)
  File "/Users/jessegrabowski/mambaforge/envs/pytensor-dev/lib/python3.12/site-packages/IPython/core/async_helpers.py", line 128, in _pseudo_sync_runner
    coro.send(None)
  File "/Users/jessegrabowski/mambaforge/envs/pytensor-dev/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3362, in run_cell_async
    has_raised = await self.run_ast_nodes(code_ast.body, cell_name,
  File "/Users/jessegrabowski/mambaforge/envs/pytensor-dev/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3607, in run_ast_nodes
    if await self.run_code(code, result, async_=asy):
  File "/Users/jessegrabowski/mambaforge/envs/pytensor-dev/lib/python3.12/site-packages/IPython/core/interactiveshell.py", line 3667, in run_code
    exec(code_obj, self.user_global_ns, self.user_ns)
  File "/var/folders/7b/rzxy96cj0w751_6td3g2yss00000gn/T/ipykernel_46141/4000051598.py", line 7, in <module>
    z = x @ y

HINT: Use the PyTensor flag `exception_verbosity=high` for a debug print-out and storage map footprint of this Apply node.

Contributor guide

Open the contributing guide

First steps

  1. Read the whole issue, then the project's contributing guide.
  2. Comment on the issue to say you are picking it up — it saves two people doing the same work.
  3. Fork the repository and make your change on a branch.
  4. Open a pull request that references the issue number.

Research direction

Start by running the Python and NumPy reproduction in the issue and inspect the BatchedDot entry point that raises the error. Trace how its inputs are handled for complex128 values, then add coverage for the shown case. Done means batched dot evaluates successfully for complex inputs without the reported NotImplementedError.

Written by the indexing model from the issue text.

Assessment

Tech stack
numpy, python
Domain
machine-learning
Issue type
Bug
Difficulty
3/5
Estimated time
1-2 days
Activity status
Stale
Clarity
Mostly clear
Newbie friendliness
43/100

Get new issues in your inbox

A short digest of beginner-friendly GitHub issues.