guard_negative_gather_indices: offset gather_nd indices from batch_dims

[LeSingh1]

Problem

gather_nd's indices address x starting at dimension batch_dims — its documented output is V = K[:-1] + D[batch_dims + K[-1]:]. So a negative index at position j of the last axis must be offset by D[batch_dims + j].

guard_negative_gather_indices always slices x's shape from 0:

slice_shape = mb.slice_by_size(x=x_shape, begin=[0], size=indices_last_dim_expand)
indices_plus = mb.add(x=indices_int32, y=slice_shape)

With batch_dims > 0 that adds the wrong dimension size:

x_val   = np.arange(2 * 5 * 3).reshape(2, 5, 3).astype(np.float32)   # (2, 5, 3)
indices = np.array([[[-1], [0]], [[-1], [1]]], dtype=np.int32)       # (2, 2, 1)

@mb.program(input_specs=[], opset_version=ct.target.iOS17)
def prog():
    return mb.gather_nd(x=x_val, indices=indices, batch_dims=1)

PASS_REGISTRY["common::guard_negative_gather_indices"](prog)

The rewrite is fully constant-foldable here, so the wrong result is directly visible:

	value
slice_by_size begin	[0] — should be [1]
rewritten indices	[[[1], [0]], [[1], [1]]]
correct	[[[4], [0]], [[4], [1]]]

-1 should become -1 + x.shape[batch_dims] = -1 + 5 = 4; the pass produces -1 + x.shape[0] = 1. The model gathers a different element, and goes out of range entirely whenever x.shape[0] < -min(indices).

This pass runs for iOS17+ targets, where gather_nd no longer accepts negative indices — so this is exactly the path meant to make negative indices safe.

Fix

Slice x's shape from batch_dims instead of 0. batch_dims == 0 is unchanged.

The gather branch is already correct and is left alone: gather's axis is absolute with respect to x and its indices are documented as -D[axis] <= v < D[axis], so x.shape[axis] is the right offset whatever batch_dims is.

Tests

TestGuardNegativeGatherIndices::test_guard_negative_gather_nd_indices_with_batch_dims[batch_dims=0,1] — asserts the slice_by_size begin and the constant-folded rewritten indices. batch_dims=1 fails on main; batch_dims=0 passes both before and after, guarding against over-correction.

The existing test_guard_negative_gather_indices (which only uses the default batch_dims=0) is untouched. I ran the whole class before and after; the pre-existing failure set is identical (this environment cannot load CoreML.framework, so assert_model_is_valid fails there either way).

Note: this touches optimize_tensor_operation.py and test_passes.py, which my PR #2789 also modifies, in a different pass/class.