keras/src/backend/torch/numpy.py/split
def split(x, indices_or_sections, axis=0):
    x = convert_to_tensor(x)
    dim = x.shape[axis]
    if not isinstance(indices_or_sections, int):
        indices_or_sections = convert_to_tensor(indices_or_sections)
        start_size = indices_or_sections[0:1]
        end_size = dim - indices_or_sections[-1:]
        chunk_sizes = torch.concat(
            [start_size, torch.diff(indices_or_sections), end_size], dim=0
        )
        # torch.split doesn't support tensor input for `split_size_or_sections`
        chunk_sizes = chunk_sizes.tolist()
    else:
        if dim % indices_or_sections != 0:
            raise ValueError(
                f"Received indices_or_sections={indices_or_sections} "
                f"(interpreted as a number of sections) and axis={axis}, "
                f"but input dimension x.shape[{axis}]={x.shape[axis]} "
                f"is not divisible by {indices_or_sections}. "
                f"Full input shape: x.shape={x.shape}"
            )
        chunk_sizes = dim // indices_or_sections
    out = torch.split(
        tensor=x,
        split_size_or_sections=chunk_sizes,
        dim=axis,
    )
    if dim == 0 and isinstance(indices_or_sections, int):
        out = [out[0].clone() for _ in range(indices_or_sections)]
    return list(out)