Check failed in jt.split/jt.unbind/jt.nn.bilinear/jt.nn.grid_sampler_2d/jt.size

[x0w3n]

Describe the bug

These apis trigger Check failed if the dimensions do not match during processing, and the trigger reason is similar to that of log.

Full Log

[i 1203 16:06:45.863062 20 compiler.py:956] Jittor(1.3.9.10) src: /home/miniconda3/envs/jittor/lib/python3.9/site-packages/jittor
[i 1203 16:06:45.871026 20 compiler.py:957] g++ at /usr/bin/g++(11.2.0)
[i 1203 16:06:45.871104 20 compiler.py:958] cache_path: /home/.cache/jittor/jt1.3.9/g++11.2.0/py3.9.12/Linux-5.15.0-1x30/INTELRXEONRGOLx51/ef26/default
[i 1203 16:06:45.957956 20 install_cuda.py:93] cuda_driver_version: [12, 2]
[i 1203 16:06:45.972899 20 __init__.py:412] Found /home/.cache/jittor/jtcuda/cuda12.2_cudnn8_linux/bin/nvcc(12.2.140) at /home/.cache/jittor/jtcuda/cuda12.2_cudnn8_linux/bin/nvcc.
[i 1203 16:06:46.427616 20 __init__.py:412] Found gdb(12.0.90) at /usr/bin/gdb.
[i 1203 16:06:46.430119 20 __init__.py:412] Found addr2line(2.38) at /usr/bin/addr2line.
[i 1203 16:06:46.851160 20 compiler.py:1013] cuda key:cu12.2.140_sm_89
[i 1203 16:06:48.011471 20 __init__.py:227] Total mem: 251.50GB, using 16 procs for compiling.
[i 1203 16:06:48.232359 20 jit_compiler.cc:28] Load cc_path: /usr/bin/g++
[i 1203 16:06:48.448548 20 init.cc:63] Found cuda archs: [89,]
Traceback (most recent call last):
  File "test.py", line 5, in <module>
    out_parts = jt.split(data, split_size=1, dim=6)
  File "/home/miniconda3/envs/jittor/lib/python3.9/site-packages/jittor/misc.py", line 880, in split
    shape = d.shape[dim]
RuntimeError: Wrong inputs arguments, Please refer to examples(help(jt.NanoVector.__map_getitem__)).

Types of your inputs are:
 self   = NanoVector,
 arg0   = int,

The function declarations are:
 inline int64 at(int i)
 inline NanoVector slice(Slice slice)

Failed reason:[f 1203 16:06:49.362037 20 nano_vector.h:116] Check failed: i>=0 && i<size()  Something wrong... Could you please report this issue?

Minimal Reproduce

jt.split

import jittor as jt

data = jt.ones((5, 5))
out_parts = jt.split(data, split_size=1, dim=6)
print("Split output parts:", out_parts)

jt.unbind

import jittor as jt

a = jt.random((2,3))
b = jt.unbind(a, 6)

jt.nn.bilinear

import jittor as jt

in1 = jt.ones((10, 5)) 
in2 = jt.ones(5)      
weight = jt.ones((7, 5)) 
bias = jt.ones((5,))   

result = jt.nn.bilinear(in1, in2, weight, bias)
print(result)

jt.nn.grid_sampler_2d

import jittor as jt
from jittor.nn import grid_sampler_2d
grid = jt.array([[[[0, 0], [1, 1]], [[1, 0], [0, 1]]]])
X = jt.array([[1, 2], [3, 4]])
mode = 'bilinear'
padding_mode = 'border'
align_corners = True

sampled_X = grid_sampler_2d(X, grid, mode, padding_mode, align_corners)

jt.size

import jittor as jt

x = jt.rand((3, 3))
print(jt.size(x, 3))

Expected behavior