code: def test_local_Unique_Alloc_lift(
      x_val,
      axis,
      new_shape,
      return_index,
      return_counts,
      return_inverse,
  ):
      x = as_tensor_variable(x_val).type()
      y = unique(
          alloc(x, *new_shape),
          return_index=return_index,
          return_counts=return_counts,
          return_inverse=return_inverse,
          axis=axis,
      )
  
      if isinstance(y, list):
          y, *_ = y
  
      # This approach allows us to directly confirm that `x` is in the result.
      y_fg = FunctionGraph(outputs=[y], copy_inputs=False)
      y_rewritten_fg = rewrite_graph(
          y_fg,
          clone=False,
          include=["canonicalize", "local_Unique_Alloc_lift"],
          exclude=["local_Unique_scalar"],
      )
      y_rewritten = y_rewritten_fg.outputs[0]
      y_rewritten_start = y_rewritten
  
      assert isinstance(y_rewritten_start.owner.op, Unique)
      assert y_rewritten_start.owner.inputs[0] == x
      assert not any(
          isinstance(node.op, Alloc)
          for node in y_rewritten_fg.apply_nodes
      )
  
      default_mode = get_default_mode()
      # The rewrite has already been applied to `y_rewritten`, so we can--and
      # should--exclude it from the compilation of both our reference, `y`, and
      # the rewritten result, `y_rewritten`.
      # The remaining exclusions simply allow us to perform the check below that
      # makes sure the original `Alloc` is present in our reference (sub)graph.
      rewrite_mode = default_mode.excluding(
          "local_useless_alloc",
          "local_alloc_sink_dimshuffle",
          "local_Unique_Alloc_lift",
      )
      y_fn = function([x], [y, y_rewritten], mode=rewrite_mode)
      # Make sure that the original `Alloc` is used to compute the reference `y`
      # result
      assert any(
          isinstance(node.op, Alloc)
          for node in y_fn.maker.fgraph.apply_nodes
      )
  
      y_exp_val, y_val = y_fn(x_val)
      assert np.array_equal(y_exp_val, y_val)
  
file_path: pytensor/tests/tensor/rewriting/test_extra_ops.py
imports: import numpy
  import pytest
  import pytensor.scalar
  from pytensor.compile.function import function
  from pytensor.compile.mode import OPT_NONE, Mode, get_default_mode
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.rewriting.utils import rewrite_graph
  from pytensor.tensor.basic import Alloc, alloc, as_tensor_variable, second
  from pytensor.tensor.elemwise import DimShuffle, Elemwise
  from pytensor.tensor.extra_ops import Repeat, Unique, repeat, unique
  from pytensor.tensor.type import dscalar
label_desc: Use image node_id 2 for calling a global function with example usage: test_local_Unique_Alloc_lift(x_val, axis, new_shape, return_index, return_counts, return_inverse) without return types