code: def test_basic(self):
      a = matrix("a")
      a_splits = ptb.split(a, splits_size=[1, 2, 3], n_splits=3, axis=0)
      fg = FunctionGraph([a], a_splits)
      compare_jax_and_py(
          fg,
          [
              np.zeros((6, 4)).astype(config.floatX),
          ],
      )
  
      a = matrix("a", shape=(6, None))
      a_splits = ptb.split(
          a, splits_size=[2, a.shape[0] - 2], n_splits=2, axis=0
      )
      fg = FunctionGraph([a], a_splits)
      compare_jax_and_py(
          fg,
          [
              np.zeros((6, 4)).astype(config.floatX),
          ],
      )
  
imports: import numpy
  import pytest
  from pytensor.compile import get_mode
  import jax.errors
  import pytensor
  import pytensor.tensor.basic
  from pytensor.configdefaults import config
  from pytensor.graph.fg import FunctionGraph
  from pytensor.graph.op import get_test_value
  from pytensor.tensor.type import iscalar, matrix, scalar, vector
  from tests.link.jax.test_basic import compare_jax_and_py
  from tests.tensor.test_basic import TestAlloc
label_desc: Use image node_id 1 for calling the TestJaxSplit obj's underlying member method code with example usage: obj.test_basic() without return types