code: def test_jax_block_diag_blockwise():
      A = pt.tensor3("A")
      B = pt.tensor3("B")
      out = pt_slinalg.block_diag(A, B)
      out_fg = FunctionGraph([A, B], [out])
      compare_jax_and_py(
          out_fg,
          [
              np.random.normal(size=(5, 5, 5)).astype(config.floatX),
              np.random.normal(size=(5, 3, 3)).astype(config.floatX),
          ],
      )
  
file_path: pytensor/tests/link/jax/test_slinalg.py
imports: import numpy
  import pytest
  import pytensor.tensor
  from pytensor.configdefaults import config
  from pytensor.graph.fg import FunctionGraph
  from pytensor.tensor import nlinalg
  from pytensor.tensor import slinalg
  from pytensor.tensor import subtensor
  from pytensor.tensor.math import clip, cosh
  from pytensor.tensor.type import matrix, vector
  from tests.link.jax.test_basic import compare_jax_and_py
label_desc: Use image node_id 4 for calling a global function with example usage: test_jax_block_diag_blockwise() without return types