Modify np.tri Op to use _iota instead

[Nimish-4]

Description

Related Issue

• Closes #
• Related to Replace our Tri op with an OpFromGraph #1265

Checklist

• Checked that the pre-commit linting/style checks pass
• Included tests that prove the fix is effective or that the new feature works
• Added necessary documentation (docstrings and/or example notebooks)
• If you are a pro: each commit corresponds to a relevant logical change

Type of change

• New feature / enhancement
• Bug fix
• Documentation
• Maintenance
• Other (please specify):

---

📚 Documentation preview 📚: https://pytensor--1276.org.readthedocs.build/en/1276/

[jessegrabowski]

The circular import is because you are importing _iota from tensor.einsum inside tensor.basic, but inside tensor.einsum, tensor.basic is imported.

The solution is to move the _iota function to tensor.basic. I also suggest to remove the leading underscore, because there's no reason this function should be considered "hidden". I don't think it needs to be added to __all__ (because it's not a function that exists in numpy), but I also wouldn't object to it.

[ricardoV94]

Those Ops and a few more (arange, alloc, ...) should probably be in a tensor_creation file. basic is hosting too much

[Nimish-4]

@jessegrabowski My bad for missing it, I'll make the change.
@ricardoV94 If that is appropriate in scope for me, I'll be okay with working on it.

[ricardoV94]

@ricardoV94 If that is appropriate in scope for me, I'll be okay with working on it.

That's up to you. It would be much appreciated!

[Nimish-4]

@jessegrabowski Made the change. I have left the docstring of iota mostly the same, but I guess the example section may need to be removed now? Tests are pending, I'll try adding some and see how it goes.

[Nimish-4]

@jessegrabowski The tests were passing for tri but now the tests for tril and triu are failing. The issue seems, in the test we are passing the matrix as a symbolic variable, and the call to Tri here also passes M and dtype as symbolic variables (because they are derived from the matrix?) which is not allowed according to the current structure. Any suggestions on how to move ahead? Still figuring out symbolic variables.

[jessegrabowski]

Tri doesn't need init at all, just make it an empty OpFromGraph, then make a function called tri that actually does the work. Have a look at how the kronecker product is implemented for a template:

[Nimish-4]

@jessegrabowski I made the change, the tests are passing now. I have removed 'complex64' dtype from the tests for now, because there seemed to be some compilation error which I could not remove.

[jessegrabowski]

Well what was the error? I don't want us removing tests.

[Nimish-4]

@jessegrabowski The stack trace looks like (removed some parts)

[jessegrabowski]

Put the test back and i'll trigger a CI run so I can see the full output

[Nimish-4]

@jessegrabowski Also stuck at this test case; seems to be failing because we are passing concrete values instead of symbolic ones (I think). The same is tested in the test case below it.

[jessegrabowski]

Suggested change

	M = N
	M = N.copy()

OpFromGraph can freak out sometimes if a single input is re-used

[Nimish-4]

@jessegrabowski Makes sense, I'll make the change. But it still resulted in the same error.

Stack trace (formatting may be unclear)

tests\link\jax\test_tensor_basic.py F [100%]

====================================================== FAILURES =======================================================
______________________________________________________ test_tri _______________________________________________________

def streamline_default_f():
    for x in no_recycling:
        x[0] = None
    try:
        # strict=False because we are in a hot loop
        for thunk, node, old_storage in zip(
            thunks, order, post_thunk_old_storage, strict=False
        ):

          thunk()

pytensor\link\utils.py:197:

---

pytensor\graph\op.py:531: in rval
r = p(n, [x[0] for x in i], o)
pytensor\compile\builders.py:875: in perform
variables = self.fn(*inputs)
pytensor\compile\builders.py:856: in fn
self.fn = function(self.inner_inputs, self.inner_outputs, **self.kwargs)
pytensor\compile\function_init.py:332: in function
fn = pfunc(
pytensor\compile\function\pfunc.py:472: in pfunc
return orig_function(
pytensor\compile\function\types.py:1820: in orig_function
m = Maker(
pytensor\compile\function\types.py:1567: in init
self.check_unused_inputs(inputs, outputs, on_unused_input)

---

inputs = [In(*0-<Scalar(int8, shape=())>), In(*1-<Scalar(int8, shape=())>), In(*2-<Scalar(int8, shape=())>)]
outputs = [Out(Cast{float64}.0,False)], on_unused_input = 'raise'

@staticmethod
def check_unused_inputs(inputs, outputs, on_unused_input):
    if on_unused_input is None:
        on_unused_input = config.on_unused_input

    if on_unused_input == "ignore":
        return

    # There should be two categories of variables in inputs:
    #  - variables that have to be provided (used_inputs)
    #  - shared variables that will be updated
    used_inputs = list(
        ancestors(
            (
                [o.variable for o in outputs]
                + [
                    i.update
                    for i in inputs
                    if getattr(i, "update", None) is not None
                ]
            ),
            blockers=[i.variable for i in inputs],
        )
    )

    msg = (
        "pytensor.function was asked to create a function computing "
        "outputs given certain inputs, but the provided input "
        "variable at index %i is not part of the computational graph "
        "needed to compute the outputs: %s.\n%s"
    )
    warn_msg = (
        "To make this warning into an error, you can pass the "
        "parameter on_unused_input='raise' to pytensor.function. "
        "To disable it completely, use on_unused_input='ignore'."
    )
    err_msg = (
        "To make this error into a warning, you can pass the "
        "parameter on_unused_input='warn' to pytensor.function. "
        "To disable it completely, use on_unused_input='ignore'."
    )

    for i in inputs:
        if (i.variable not in used_inputs) and (i.update is None):
            if on_unused_input == "warn":
                warnings.warn(
                    msg % (inputs.index(i), i.variable, warn_msg), stacklevel=6
                )
            elif on_unused_input == "raise":

              raise UnusedInputError(msg % (inputs.index(i), i.variable, err_msg))

E pytensor.compile.function.types.UnusedInputError: pytensor.function was asked to create a function computing outputs given certain inputs, but the provided input variable at index 0 is not part of the computational graph needed to compute the outputs: *0-<Scalar(int8, shape=())>.
E To make this error into a warning, you can pass the parameter on_unused_input='warn' to pytensor.function. To disable it completely, use on_unused_input='ignore'.

pytensor\compile\function\types.py:1438: UnusedInputError

During handling of the above exception, another exception occurred:

def test_tri():
    out = ptb.tri(10, 10, 0)

  compare_jax_and_py([], [out], [])

tests\link\jax\test_tensor_basic.py:207:

---

tests\link\jax\test_basic.py:87: in compare_jax_and_py
py_res = pytensor_py_fn(*test_inputs)
pytensor\compile\function\types.py:1037: in call
outputs = vm() if output_subset is None else vm(output_subset=output_subset)
pytensor\link\utils.py:201: in streamline_default_f
raise_with_op(fgraph, node, thunk)
pytensor\link\utils.py:526: in raise_with_op
raise exc_value.with_traceback(exc_trace)
pytensor\link\utils.py:197: in streamline_default_f
thunk()
pytensor\graph\op.py:531: in rval
r = p(n, [x[0] for x in i], o)
pytensor\compile\builders.py:875: in perform
variables = self.fn(*inputs)
pytensor\compile\builders.py:856: in fn
self.fn = function(self.inner_inputs, self.inner_outputs, **self.kwargs)
pytensor\compile\function_init.py:332: in function
fn = pfunc(
pytensor\compile\function\pfunc.py:472: in pfunc
return orig_function(
pytensor\compile\function\types.py:1820: in orig_function
m = Maker(
pytensor\compile\function\types.py:1567: in init
self.check_unused_inputs(inputs, outputs, on_unused_input)

---

inputs = [In(*0-<Scalar(int8, shape=())>), In(*1-<Scalar(int8, shape=())>), In(*2-<Scalar(int8, shape=())>)]
outputs = [Out(Cast{float64}.0,False)], on_unused_input = 'raise'

@staticmethod
def check_unused_inputs(inputs, outputs, on_unused_input):
    if on_unused_input is None:
        on_unused_input = config.on_unused_input

    if on_unused_input == "ignore":
        return

    # There should be two categories of variables in inputs:
    #  - variables that have to be provided (used_inputs)
    #  - shared variables that will be updated
    used_inputs = list(
        ancestors(
            (
                [o.variable for o in outputs]
                + [
                    i.update
                    for i in inputs
                    if getattr(i, "update", None) is not None
                ]
            ),
            blockers=[i.variable for i in inputs],
        )
    )

    msg = (
        "pytensor.function was asked to create a function computing "
        "outputs given certain inputs, but the provided input "
        "variable at index %i is not part of the computational graph "
        "needed to compute the outputs: %s.\n%s"
    )
    warn_msg = (
        "To make this warning into an error, you can pass the "
        "parameter on_unused_input='raise' to pytensor.function. "
        "To disable it completely, use on_unused_input='ignore'."
    )
    err_msg = (
        "To make this error into a warning, you can pass the "
        "parameter on_unused_input='warn' to pytensor.function. "
        "To disable it completely, use on_unused_input='ignore'."
    )

    for i in inputs:
        if (i.variable not in used_inputs) and (i.update is None):
            if on_unused_input == "warn":
                warnings.warn(
                    msg % (inputs.index(i), i.variable, warn_msg), stacklevel=6
                )
            elif on_unused_input == "raise":

              raise UnusedInputError(msg % (inputs.index(i), i.variable, err_msg))

E pytensor.compile.function.types.UnusedInputError: pytensor.function was asked to create a function computing outputs given certain inputs, but the provided input variable at index 0 is not part of the computational graph needed to compute the outputs: *0-<Scalar(int8, shape=())>.
E To make this error into a warning, you can pass the parameter on_unused_input='warn' to pytensor.function. To disable it completely, use on_unused_input='ignore'.
E Apply node that caused the error: Tri{inline=False}(10, 10, 0)
E Toposort index: 0
E Inputs types: [TensorType(int8, shape=()), TensorType(int8, shape=()), TensorType(int8, shape=())]
E Inputs shapes: [(), (), ()]
E Inputs strides: [(), (), ()]
E Inputs values: [array(10, dtype=int8), array(10, dtype=int8), array(0, dtype=int8)]
E Outputs clients: [[output0]]
E
E Backtrace when the node is created (use PyTensor flag traceback__limit=N to make it longer):
E File "C:\Users\Public\miniforge3\envs\pytensor-dev\Lib\site-packages\pluggy_callers.py", line 103, in _multicall
E res = hook_impl.function(*args)
E File "C:\Users\Public\miniforge3\envs\pytensor-dev\Lib\site-packages_pytest\runner.py", line 174, in pytest_runtest_call
E item.runtest()
E File "C:\Users\Public\miniforge3\envs\pytensor-dev\Lib\site-packages_pytest\python.py", line 1627, in runtest
E self.ihook.pytest_pyfunc_call(pyfuncitem=self)
E File "C:\Users\Public\miniforge3\envs\pytensor-dev\Lib\site-packages\pluggy_hooks.py", line 513, in call
E return self._hookexec(self.name, self._hookimpls.copy(), kwargs, firstresult)
E File "C:\Users\Public\miniforge3\envs\pytensor-dev\Lib\site-packages\pluggy_manager.py", line 120, in _hookexec
E return self._inner_hookexec(hook_name, methods, kwargs, firstresult)
E File "C:\Users\Public\miniforge3\envs\pytensor-dev\Lib\site-packages\pluggy_callers.py", line 103, in _multicall
E res = hook_impl.function(*args)
E File "C:\Users\Public\miniforge3\envs\pytensor-dev\Lib\site-packages_pytest\python.py", line 159, in pytest_pyfunc_call
E result = testfunction(**testargs)
E File "C:\Users\Nimish Purohit\pytensor\tests\link\jax\test_tensor_basic.py", line 205, in test_tri
E out = ptb.tri(10, 10, 0)
E
E HINT: Use the PyTensor flag exception_verbosity=high for a debug print-out and storage map footprint of this Apply node.

pytensor\compile\function\types.py:1438: UnusedInputError
================================================ slowest 50 durations =================================================
0.22s call tests/link/jax/test_tensor_basic.py::test_tri

(2 durations < 0.005s hidden. Use -vv to show these durations.)
=============================================== short test summary info ===============================================
FAILED tests/link/jax/test_tensor_basic.py::test_tri - pytensor.compile.function.types.UnusedInputError: pytensor.function was asked to create a function computing output...
================================================== 1 failed in 5.51s ==================================================

[jessegrabowski]

I'll have a look more closely at what is going on over the weekend. It's not obvious to me at first glance

[Nimish-4]

@jessegrabowski Any idea on how to proceed? The test function call will result in this call, which is causing the UnusedInputError. This call will contain an empty graph_inputs, compared to the normal python test cases that are passing.