Check failed in jt.misc.chunk/jt.misc.cross/jt.misc.kthvalue

[x0w3n]

Describe the bug

Check failed is triggered when the dim of jt.misc.chunk/jt.misc.cross/jt.misc.kthvalue exceeds the dimensions of the input tensor.
The error log suggests reporting this issue.

Full Log

jt.misc.chunk

[i 1203 10:45:13.413438 48 compiler.py:956] Jittor(1.3.9.10) src: /home/miniconda3/envs/jittor/lib/python3.9/site-packages/jittor
[i 1203 10:45:13.424089 48 compiler.py:957] g++ at /usr/bin/g++(11.2.0)
[i 1203 10:45:13.424608 48 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 10:45:13.498549 48 install_cuda.py:93] cuda_driver_version: [12, 2]
[i 1203 10:45:13.513192 48 __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 10:45:13.966581 48 __init__.py:412] Found gdb(12.0.90) at /usr/bin/gdb.
[i 1203 10:45:13.969077 48 __init__.py:412] Found addr2line(2.38) at /usr/bin/addr2line.
[i 1203 10:45:14.196975 48 compiler.py:1013] cuda key:cu12.2.140_sm_89
[i 1203 10:45:15.333428 48 __init__.py:227] Total mem: 251.50GB, using 16 procs for compiling.
[i 1203 10:45:15.500596 48 jit_compiler.cc:28] Load cc_path: /usr/bin/g++
[i 1203 10:45:15.731666 48 init.cc:63] Found cuda archs: [89,]
Traceback (most recent call last):
  File "test.py", line 5, in <module>
    chunks = jt.misc.chunk(t, chunks=1, dim=4)
  File "/home/miniconda3/envs/jittor/lib/python3.9/site-packages/jittor/misc.py", line 263, in chunk
    l = x.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 10:45:16.396084 48 nano_vector.h:116] Check failed: i>=0 && i<size()  Something wrong... Could you please report this issue?

jt.misc.cross

[i 1203 10:47:27.925719 36 compiler.py:956] Jittor(1.3.9.10) src: /home/miniconda3/envs/jittor/lib/python3.9/site-packages/jittor
[i 1203 10:47:27.935346 36 compiler.py:957] g++ at /usr/bin/g++(11.2.0)
[i 1203 10:47:27.935864 36 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 10:47:28.006093 36 install_cuda.py:93] cuda_driver_version: [12, 2]
[i 1203 10:47:28.021448 36 __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 10:47:28.289650 36 __init__.py:412] Found gdb(12.0.90) at /usr/bin/gdb.
[i 1203 10:47:28.292175 36 __init__.py:412] Found addr2line(2.38) at /usr/bin/addr2line.
[i 1203 10:47:28.618039 36 compiler.py:1013] cuda key:cu12.2.140_sm_89
[i 1203 10:47:29.828549 36 __init__.py:227] Total mem: 251.50GB, using 16 procs for compiling.
[i 1203 10:47:30.048711 36 jit_compiler.cc:28] Load cc_path: /usr/bin/g++
[i 1203 10:47:30.309785 36 init.cc:63] Found cuda archs: [89,]
Traceback (most recent call last):
  File "test.py", line 5, in <module>
    result = jt.misc.cross(x, x, dim=4)
  File "/home/miniconda3/envs/jittor/lib/python3.9/site-packages/jittor/misc.py", line 437, in cross
    assert input.shape[dim] == 3, "input dim shape must be 3"
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 10:47:31.011993 36 nano_vector.h:116] Check failed: i>=0 && i<size()  Something wrong... Could you please report this issue?

jt.misc.kthvalue

[i 1203 10:58:59.449596 96 compiler.py:956] Jittor(1.3.9.10) src: /home/miniconda3/envs/jittor/lib/python3.9/site-packages/jittor
[i 1203 10:58:59.461140 96 compiler.py:957] g++ at /usr/bin/g++(11.2.0)
[i 1203 10:58:59.461609 96 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 10:58:59.528658 96 install_cuda.py:93] cuda_driver_version: [12, 2]
[i 1203 10:58:59.540738 96 __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 10:58:59.856677 96 __init__.py:412] Found gdb(12.0.90) at /usr/bin/gdb.
[i 1203 10:58:59.859335 96 __init__.py:412] Found addr2line(2.38) at /usr/bin/addr2line.
[i 1203 10:59:00.377635 96 compiler.py:1013] cuda key:cu12.2.140_sm_89
[i 1203 10:59:01.620442 96 __init__.py:227] Total mem: 251.50GB, using 16 procs for compiling.
[i 1203 10:59:01.817205 96 jit_compiler.cc:28] Load cc_path: /usr/bin/g++
[i 1203 10:59:02.028876 96 init.cc:63] Found cuda archs: [89,]
Traceback (most recent call last):
  File "test.py", line 7, in <module>
    result, indices = jt.misc.kthvalue(input_var, k, dim)
  File "/home/miniconda3/envs/jittor/lib/python3.9/site-packages/jittor/misc.py", line 1002, in kthvalue
    index,values = jt.argsort(input,dim=dim)
RuntimeError: Wrong inputs arguments, Please refer to examples(help(jt.ops.argsort)).

Types of your inputs are:
 self   = module,
 args   = (Var, ),
 kwargs = {dim=int, },

The function declarations are:
 vector_to_tuple<VarHolder*> argsort(VarHolder* x,  int dim=-1,  bool descending=false,  NanoString dtype=ns_int32)

Failed reason:[f 1203 10:59:02.901706 96 argsort_op.cc:127] Check failed dim(1) < (int)x->shape.size()(1) Something wrong ... Could you please report this issue?

Minimal Reproduce

import jittor as jt

t = jt.rand((5,))
chunks = jt.misc.chunk(t, chunks=1, dim=4)
chunks

import jittor as jt

x = jt.rand(2, 3)
result = jt.misc.cross(x, x, dim=4)
print(result)

import jittor as jt

input_var = jt.array([1, 2, 3])
k = 1 
dim = 1
result, indices = jt.misc.kthvalue(input_var, k, dim)
print("Result:", result)
print("Indices:", indices)

Expected behavior