code: def local_blockwise_alloc(fgraph, node):
      """Push Allocs from the inputs to the output of Blockwise Ops.
  
      BOp = Blockwise(Op, signature="(x),(x)->(x)")
      BOp(vector, alloc(vector, 10, 5)) -> alloc(BOp)(vector, vector), 10, 5)
      BOp(vector, alloc(scalar, 10, 5)) -> alloc(BOp)(vector, alloc(scalar, 5), 10, 5)
      BOp(matrix, alloc(vector, 10, 5)) -> BOp(matrix, vector)
      """
  
      if not any(
          isinstance(inp.owner.op, Alloc)
          for inp in node.inputs
          if inp.owner
      ):
          return None
  
      op: Blockwise = node.op  # type: ignore
  
      batch_ndim = op.batch_ndim(node)
      if not batch_ndim:
          return None
  
      new_inputs = []
      batch_shapes = []
      can_push_any_alloc = False
      for inp, inp_sig in zip(node.inputs, op.inputs_sig):
          if inp.owner and isinstance(inp.owner.op, Alloc):
              # Push batch dims from Alloc
              value, *shape = inp.owner.inputs
  
              # Check what to do with the value of the Alloc
              squeezed_value = _squeeze_left(value, batch_ndim)
              missing_ndim = len(shape) - value.type.ndim
              if (
                  ((1,) * missing_ndim + value.type.broadcastable)[
                      batch_ndim:
                  ]
              ) != inp.type.broadcastable[batch_ndim:]:
                  # We still need an Alloc for the core dims
                  core_shape = shape[batch_ndim:]
                  # And the batch dims of the squeezed value
                  squeezed_value_batch_ndim = (
                      squeezed_value.type.ndim - len(core_shape)
                  )
                  batch_shape = [
                      1 if broadcastable else dim
                      for broadcastable, dim in zip(
                          squeezed_value.type.broadcastable[
                              :squeezed_value_batch_ndim
                          ],
                          tuple(squeezed_value.shape)[
                              :squeezed_value_batch_ndim
                          ],
                      )
                  ]
                  squeezed_value = alloc(
                      squeezed_value, *batch_shape, *core_shape
                  )
                  if (
                      squeezed_value.type.broadcastable
                      == inp.type.broadcastable
                  ):
                      # We can't change anything about this Alloc input
                      new_inputs.append(inp)
                      continue
  
              # We can push batch dims of this Alloc input
              batch_shapes.append(
                  tuple(
                      1 if broadcastable else dim
                      for broadcastable, dim in zip(
                          inp.type.broadcastable, shape[:batch_ndim]
                      )
                  )
              )
              new_inputs.append(squeezed_value)
              can_push_any_alloc = True
  
          else:
              # Nothing to do with this input other than removing dummy batch dims
              new_inputs.append(_squeeze_left(inp, batch_ndim))
  
      if not can_push_any_alloc:
          return None
  
      new_outs = node.op.make_node(*new_inputs).outputs
  
      new_out_type = new_outs[0].type
      old_out_type = node.outputs[0].type
      if new_out_type.broadcastable != old_out_type.broadcastable:
          # An Alloc is still needed to broadcast the new output to the original shape
          # We pick the most parsimonious batch dim from the pushed Alloc
          missing_ndim = old_out_type.ndim - new_out_type.ndim
          batch_shape = ([1] * missing_ndim + list(new_outs[0].shape))[
              :batch_ndim
          ]
          for i, batch_dims in enumerate(
              zip(*batch_shapes)
          ):  # Transpose shape tuples
              for batch_dim in batch_dims:
                  if batch_dim == 1:
                      continue
                  if isinstance(batch_dim, Constant):
                      # Give preference to Constants
                      batch_shape[i] = batch_dim
                      break
                  elif old_out_type.broadcastable[i]:
                      # Only use non Constant shapes if absolutely necessary
                      # Otherwise, we use the shape of the non-alloc output
                      batch_shape[i] = batch_dim
  
          copy_stack_trace(node.outputs, new_outs)
          new_outs = [
              alloc(
                  new_out,
                  *batch_shape,
                  *tuple(new_out.shape)[batch_ndim - missing_ndim :],
              )
              for new_out in new_outs
          ]
      assert (
          new_outs[0].type.broadcastable == old_out_type.broadcastable
      )
      copy_stack_trace(node.outputs, new_outs)
      return new_outs
  
file_path: pytensor/pytensor/tensor/rewriting/blockwise.py
imports: from typing import Optional
  from pytensor.compile.mode import optdb
  from pytensor.graph import Constant, node_rewriter
  from pytensor.graph.replace import vectorize_node
  from pytensor.graph.rewriting.basic import copy_stack_trace, out2in
  from pytensor.tensor.basic import Alloc, ARange, alloc, shape_padleft
  from pytensor.tensor.blockwise import Blockwise
  from pytensor.tensor.math import Dot
  from pytensor.tensor.rewriting.basic import register_canonicalize, register_specialize, register_stabilize
  from pytensor.tensor.subtensor import AdvancedIncSubtensor, AdvancedSubtensor, Subtensor
label_desc: Use image node_id 5 for calling a global function with example usage: local_blockwise_alloc(fgraph, node) and returns: new_outs, None, None, None