code: def test_blockwise_shape():
      # Single output
      inp = tensor(shape=(5, None, None))
      inp_test = np.zeros((5, 4, 3), dtype=config.floatX)
  
      # Shape can be inferred from inputs
      op = Blockwise(test_op, signature="(m, n) -> (n, m)")
      out = op(inp)
      assert out.type.shape == (5, None, None)
  
      shape_fn = pytensor.function([inp], out.shape)
      assert not any(
          isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
          for n in shape_fn.maker.fgraph.apply_nodes
      )
      assert tuple(shape_fn(inp_test)) == (5, 3, 4)
  
      # Shape can only be partially inferred from inputs
      op = Blockwise(test_op, signature="(m, n) -> (m, k)")
      out = op(inp)
      assert out.type.shape == (5, None, None)
  
      shape_fn = pytensor.function([inp], out.shape)
      assert any(
          isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
          for n in shape_fn.maker.fgraph.apply_nodes
      )
  
      shape_fn = pytensor.function([inp], out.shape[:-1])
      assert not any(
          isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
          for n in shape_fn.maker.fgraph.apply_nodes
      )
      assert tuple(shape_fn(inp_test)) == (5, 4)
  
      # Mutiple outputs
      inp1 = tensor(shape=(7, 1, None, None))
      inp2 = tensor(shape=(1, 5, None, None))
      inp1_test = np.zeros((7, 1, 4, 3), dtype=config.floatX)
      inp2_test = np.zeros((1, 5, 4, 3), dtype=config.floatX)
  
      op = Blockwise(
          test_op, signature="(m, n), (m, n) -> (n, m), (m, k)"
      )
      outs = op(inp1, inp2)
      assert outs[0].type.shape == (7, 5, None, None)
      assert outs[1].type.shape == (7, 5, None, None)
  
      shape_fn = pytensor.function(
          [inp1, inp2], [out.shape for out in outs]
      )
      assert any(
          isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
          for n in shape_fn.maker.fgraph.apply_nodes
      )
  
      shape_fn = pytensor.function([inp1, inp2], outs[0].shape)
      assert not any(
          isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
          for n in shape_fn.maker.fgraph.apply_nodes
      )
      assert tuple(shape_fn(inp1_test, inp2_test)) == (7, 5, 3, 4)
  
      shape_fn = pytensor.function(
          [inp1, inp2], [outs[0].shape, outs[1].shape[:-1]]
      )
      assert not any(
          isinstance(getattr(n.op, "core_op", n.op), MyTestOp)
          for n in shape_fn.maker.fgraph.apply_nodes
      )
      assert tuple(shape_fn(inp1_test, inp2_test)[0]) == (7, 5, 3, 4)
      assert tuple(shape_fn(inp1_test, inp2_test)[1]) == (7, 5, 4)
  
imports: import re
  from itertools import product
  from typing import Optional, Union
  import numpy
  import pytest
  import pytensor
  from pytensor import config, function
  from pytensor.compile import get_mode
  from pytensor.gradient import grad
  from pytensor.graph import Apply, Op
  from pytensor.graph.replace import vectorize_node
  from pytensor.raise_op import assert_op
  from pytensor.tensor import diagonal, log, tensor
  from pytensor.tensor.blockwise import Blockwise, vectorize_node_fallback
  from pytensor.tensor.nlinalg import MatrixInverse
  from pytensor.tensor.rewriting.blas import specialize_matmul_to_batched_dot
  from pytensor.tensor.slinalg import Cholesky, Solve, cholesky, solve_triangular
  from pytensor.tensor.utils import _parse_gufunc_signature
label_desc: Use image node_id 6 for calling a global function with example usage: test_blockwise_shape() without return types