Support empty JaggedTensor structures and validate raw constructors

[swahtz]

Fixes #89, fixes #722, fixes #802, fixes #803, fixes #805, fixes #811. Also covers the JaggedTensor launch sites of #627 (JaggedTensorIndex.cu, JCat0.cu). The rest of #627 stays open.

Summary

A JaggedTensor with zero tensors, zero elements, or an empty outer list now constructs correctly and works with every JaggedTensor operation. Before this PR, such tensors either could not be built or crashed their readers. The PR defines the structural invariants once, enforces them in the public raw constructors, and fixes the readers that assumed every outer list had a tensor.

Root causes

Issue	Symptom	Cause
#89	JaggedTensor([]) raises	The list constructors rejected empty input or called torch::cat on an empty list.
#802	Empty outer lists dropped by lshape, unbind, from_data_offsets_and_list_ids and CPU indexing	Readers derived outer lists from list ids (max + 1, or "the outer id changed") instead of num_outer_lists.
#805	lshape segfaults	list_ids.numel() == 0 bypassed the row-count check, and the reader then read past the list ids.
#803	Gsplat contributor ids segfault when no pixel has a contributor	The op's empty path built the #805 structure.
#722	Malformed offsets cause out-of-bounds reads in later ops	The raw constructors never checked offset values.
#811	CUDA jcat(dim=0) writes a wrong jidx	The kernel wrote at the in-tensor position and overran the offsets by one.

Behavior changes

• New num_outer_lists keyword. from_data_offsets_and_list_ids and from_data_indices_and_list_ids take an optional num_outer_lists, because list ids cannot express trailing empty outer lists. Without it, the count is the largest outer id + 1.
• The public raw constructors now raise ValueError on malformed structure. This applies to from_data_offsets_and_list_ids, from_data_indices_and_list_ids, from_data_and_offsets and from_data_and_indices. Malformed structure includes:
  • offsets that don't start at 0, decrease, or don't end at data.shape[0];
  • indices that are unsorted or out of range;
  • list ids whose row count doesn't match the tensor count, or whose ids are unsorted or not numbered 0, 1, …
• Integer structure tensors are cast and moved instead of rejected. For example, int64 list ids are cast to the canonical dtype, and host-built offsets move to the data's device.
• The nested lsizes constructor no longer takes total_tensors. This changes the C++ signature JaggedTensor(lsizes, total_tensors, data) to JaggedTensor(lsizes, data). The Python API is unaffected.
• An empty list holds an empty CPU float32 tensor. JaggedTensor([]) and JaggedTensor([[], []]) have no tensor to take a device, dtype or element shape from, so they use these defaults. To get an empty structure on a specific device or dtype, use from_data_offsets_and_list_ids(..., num_outer_lists=N) or the from_* lsizes builders.

Also fixed

Bugs found along the way that no open issue covered:

• set_jdata rejected zero-element nested JaggedTensors, such as ray-op outputs with zero hits.
• jsum, jmin and jmax at dim 0 treated "no elements" as "one tensor": the sum had the wrong row count, and min/max raised.
• A pickle round-trip dropped trailing empty outer lists.
• CPU slices with a step dropped empty outer lists and renumbered the rest; CUDA kept them.
• fvdb/attention.py dropped the query's list ids from its output.

Performance

Per-call time on an RTX PRO 6000 with the GPU idle. "main" is a 2026-09-29 build; the code timed here has not changed on main since 2026-08-06. Lower is better.

Call	main	This PR	Change
from_data_and_offsets, ldim 1, offsets on GPU, B=1	3.9 µs	15.7 µs	~12 µs slower: one validation launch and sync
from_data_and_offsets, ldim 1, offsets on GPU, B=64	8.0 µs	19.3 µs	~11 µs slower: same
from_data_and_offsets, ldim 1, offsets on host, B=1	3.3 µs	18.7 µs	~15 µs slower: the offsets are now copied to the data device (main left CPU offsets on CUDA data); validation itself runs on the host with no sync
from_data_and_offsets, ldim 1, offsets on host, B=64	88 µs	25 µs	~3.5× faster
from_data_and_indices, ldim 1, B=1 / B=64	139 / 185 µs	29 / 33 µs	~5× faster: the fused check replaces unique_dim
from_data_and_indices, 200k elements, 4 tensors	193 µs	36 µs	~5× faster
from_data_offsets_and_list_ids, ldim 2, 64 × 64	35 µs	27 µs	slightly faster, while also validating every list id
jflatten(dim=1), ldim 2, 64 × 64	205 µs	40 µs	~5× faster: no unique_dim
jt[i] integer index, ldim 2, 64 × 64	50 µs	42 µs	slightly faster: one binary-search kernel
jagged_like (no-sync reference)	2.1 µs	1.9 µs	unchanged

• What gets slower. Only from_data_and_offsets and from_data_offsets_and_list_ids with ldim 1 and small offsets. Main did no sync there; validation needs one launch and one sync on the GPU, or a host-to-device copy when the offsets start on the host.
• What is unchanged. Op outputs (grid, mesh, gsplat) and jagged_like skip validation through the internal unchecked constructors. Their code path is the same as main.
• Cost inside a busy training step. I also timed each call with ~2 ms of GPU work queued ahead of it. The run-to-run noise there was about ±50 µs, larger than these differences, so those numbers are not shown. Expect a GPU-side validation to add one sync's pipeline bubble when the host would otherwise run ahead.
• Reality-capture. The sparse rasterizer's backward rebuilds its pixel JaggedTensor with from_data_offsets_and_list_ids every step. Keep the pixel JaggedTensor on ctx in the sparse rasterizer backward fvdb-reality-capture#344 keeps the original on ctx instead. It doesn't depend on this PR, so it can merge first.

CI note

The unit-test jobs fail only on test_viz_shutdown.py::test_viewer_shutdown_with_gpu_buffers, with Failed to create Vulkan instance. VkResult: -9 (VK_ERROR_INCOMPATIBLE_DRIVER). That test came in with #810, whose unit-test jobs were skipped. This is the first PR to run it on the GPU runners, which have no usable Vulkan driver. It is unrelated to this change, and main will hit it too. It needs a skip when Vulkan device creation fails, or a Vulkan driver on the runners.

Test plan

./build.sh install
cd tests && pytest unit -v
cd tests && pytest unit/test_jagged_tensor.py unit/test_gaussian_splat_functional.py -v
./build.sh install gtests   # then run JaggedTensorTest via ctest

• pytest unit: 3119 passed, 10 skipped.
• TestEmptyStructures and the updated tests in test_jagged_tensor.py, run on CPU and CUDA.
• ContributingGaussianIdsEmptyTests in test_gaussian_splat_functional.py covers the rasterize_contributing_gaussian_ids_sparse segfaults when no requested pixel has a contributor #803 repro cases plus multi-camera, non-uniform and dense all-background cases.
• C++ gtests that exercise JaggedTensor: JaggedTensorTest (including new cases), PackedJaggedAccessorTest, GaussianProjectionJaggedTest, GaussianComputeSparseInfoTest, GaussianComputeNanInfMaskTest, GaussianRasterizeForwardTest and GaussianRasterizeBackwardTest. All pass.
• Each issue's reproducer, run in its own process, now behaves as the issue expected.
• black and clang-format 18 are clean.
• The full C++ gtest suite has not been run. Only the JaggedTensor-related targets above were built and run.
• The intermediate commits have not been built one by one. Only the final tree has been built and tested.

🤖 Generated with Claude Code