pymc-devs / pymc-devs/pytensor
Batched dot doesn't support complex inputs
Open
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
First steps
- Read the whole issue, then the project's contributing guide.
- Comment on the issue to say you are picking it up — it saves two people doing the same work.
- Fork the repository and make your change on a branch.
- 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