Fix Split num_outputs with an empty optional input

js/web/lib/wasm/jsep/webgpu/ops/split.ts

@@ -20,6 +20,7 @@ export interface SplitAttributes extends AttributeWithCacheKey {
   readonly axis: number;
   readonly numOutputs: number;
   readonly splitSizes: number[];
+  readonly isUnevenSplitAllowed?: boolean;
 }
 
 const validateInputs = (inputs: readonly TensorView[]): void => {
@@ -38,7 +39,12 @@ const createSplitAttributesFromInputs = (
     inputs[1].getBigInt64Array().forEach((v) => splitSizes.push(Number(v)));
     numOutputs = splitSizes.length;
   }
-  return createAttributeWithCacheKey({ numOutputs, axis: attributes.axis, splitSizes });
+  return createAttributeWithCacheKey({
+    numOutputs,
+    axis: attributes.axis,
+    splitSizes,
+    isUnevenSplitAllowed: attributes.isUnevenSplitAllowed,
+  });
 };
 
 const calculateOutputIndexImpl = (numberOfTensors: number): string => `
@@ -130,17 +136,70 @@ export const createSplitProgramInfo = (inputs: readonly TensorView[], attributes
 
 export const split = (context: ComputeContext, attributes: SplitAttributes): void => {
   validateInputs(context.inputs);
-  const updatedAttributes =
+  let updatedAttributes =
     context.inputs.length === 1 ? attributes : createSplitAttributesFromInputs(context.inputs, attributes);
-  context.compute(createSplitProgramInfo(context.inputs, updatedAttributes), { inputs: [0] });
+  if (updatedAttributes.splitSizes.length === 0) {
+    const inputShape = context.inputs[0].dims;
+    const axis = ShapeUtil.normalizeAxis(updatedAttributes.axis, inputShape.length);
+    const splitDim = inputShape[axis];
+    const numOutputs = updatedAttributes.numOutputs;
+    if (numOutputs < 1) {
+      throw new Error('numOutputs must be positive');
+    }
+    if (!updatedAttributes.isUnevenSplitAllowed && splitDim % numOutputs !== 0) {
+      throw new Error('The split dimension must be evenly divisible by the number of outputs');
+    }
+    const splitSize = Math.ceil(splitDim / numOutputs);
+    const lastSplitSize = splitDim - splitSize * (numOutputs - 1);
+    if (updatedAttributes.isUnevenSplitAllowed && lastSplitSize < 0) {
+      throw new Error('numOutputs cannot produce a negative split size');
+    }
+    const splitSizes = new Array<number>(numOutputs).fill(splitSize);
+    splitSizes[numOutputs - 1] = lastSplitSize;
+    updatedAttributes = createAttributeWithCacheKey({
+      numOutputs,
+      axis: updatedAttributes.axis,
+      splitSizes,
+      isUnevenSplitAllowed: updatedAttributes.isUnevenSplitAllowed,
+    });
+  }
+  const inputShape = context.inputs[0].dims;
+  const axis = ShapeUtil.normalizeAxis(updatedAttributes.axis, inputShape.length);
+  const inputSize = ShapeUtil.size(inputShape);
+  const outputIndices: number[] = [];
+  const nonEmptySplitSizes: number[] = [];
+  for (let i = 0; i < updatedAttributes.numOutputs; i++) {
+    const splitSize = updatedAttributes.splitSizes[i];
+    if (inputSize === 0 || splitSize === 0) {
+      const outputShape = inputShape.slice();
+      outputShape[axis] = splitSize;
+      context.output(i, outputShape);
+    } else {
+      outputIndices.push(i);
+      nonEmptySplitSizes.push(splitSize);
+    }
+  }
+  if (outputIndices.length === 0) {
+    return;
+  }
+  if (outputIndices.length !== updatedAttributes.numOutputs) {
+    updatedAttributes = createAttributeWithCacheKey({
+      numOutputs: outputIndices.length,
+      axis: updatedAttributes.axis,
+      splitSizes: nonEmptySplitSizes,
+      isUnevenSplitAllowed: updatedAttributes.isUnevenSplitAllowed,
+    });
+  }
+  context.compute(createSplitProgramInfo(context.inputs, updatedAttributes), { inputs: [0], outputs: outputIndices });
 };
 
 export const parseSplitAttributes = (attributes: Record<string, unknown>): SplitAttributes => {
   const axis = attributes.axis as number;
   const splitSizes: number[] = attributes.splitSizes as number[];
   const numOutputs = (attributes.numOutputs as number) < 0 ? splitSizes.length : (attributes.numOutputs as number);
-  if (numOutputs !== splitSizes.length) {
+  const isUnevenSplitAllowed = !!attributes.isUnevenSplitAllowed;
+  if (splitSizes.length > 0 && numOutputs !== splitSizes.length) {
     throw new Error('numOutputs and splitSizes length must be equal');
   }
-  return createAttributeWithCacheKey({ axis, numOutputs, splitSizes });
+  return createAttributeWithCacheKey({ axis, numOutputs, splitSizes, isUnevenSplitAllowed });
 };

js/web/test/data/ops/split-num-outputs.jsonc

@@ -0,0 +1,88 @@
+[
+  {
+    "name": "Split 18 with num_outputs and empty optional split input",
+    "operator": "Split",
+    "opset": { "domain": "", "version": 18 },
+    "attributes": [{ "name": "num_outputs", "data": 2, "type": "int" }],
+    "cases": [
+      {
+        "name": "T[4]",
+        "inputs": [
+          { "data": [1, 2, 3, 4], "dims": [4], "type": "float32" },
+          { "data": null, "type": "int64" }
+        ],
+        "outputs": [
+          { "data": [1, 2], "dims": [2], "type": "float32" },
+          { "data": [3, 4], "dims": [2], "type": "float32" }
+        ]
+      },
+      {
+        "name": "T[5]",
+        "inputs": [
+          { "data": [1, 2, 3, 4, 5], "dims": [5], "type": "float32" },
+          { "data": null, "type": "int64" }
+        ],
+        "outputs": [
+          { "data": [1, 2, 3], "dims": [3], "type": "float32" },
+          { "data": [4, 5], "dims": [2], "type": "float32" }
+        ]
+      },
+      {
+        "name": "T[1] with a zero-sized final output",
+        "inputs": [
+          { "data": [1], "dims": [1], "type": "float32" },
+          { "data": null, "type": "int64" }
+        ],
+        "outputs": [
+          { "data": [1], "dims": [1], "type": "float32" },
+          { "data": [], "dims": [0], "type": "float32" }
+        ]
+      }
+    ]
+  },
+  {
+    "name": "Split 18 with num_outputs and a zero-sized final output",
+    "operator": "Split",
+    "opset": { "domain": "", "version": 18 },
+    "attributes": [{ "name": "num_outputs", "data": 6, "type": "int" }],
+    "cases": [
+      {
+        "name": "T[10]",
+        "inputs": [
+          { "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], "dims": [10], "type": "float32" },
+          { "data": null, "type": "int64" }
+        ],
+        "outputs": [
+          { "data": [1, 2], "dims": [2], "type": "float32" },
+          { "data": [3, 4], "dims": [2], "type": "float32" },
+          { "data": [5, 6], "dims": [2], "type": "float32" },
+          { "data": [7, 8], "dims": [2], "type": "float32" },
+          { "data": [9, 10], "dims": [2], "type": "float32" },
+          { "data": [], "dims": [0], "type": "float32" }
+        ]
+      }
+    ]
+  },
+  {
+    "name": "Split 18 with num_outputs, negative axis and empty optional split input",
+    "operator": "Split",
+    "opset": { "domain": "", "version": 18 },
+    "attributes": [
+      { "name": "axis", "data": -1, "type": "int" },
+      { "name": "num_outputs", "data": 2, "type": "int" }
+    ],
+    "cases": [
+      {
+        "name": "T[2,5]",
+        "inputs": [
+          { "data": [1, 2, 3, 4, 5, 6, 7, 8, 9, 10], "dims": [2, 5], "type": "float32" },
+          { "data": null, "type": "int64" }
+        ],
+        "outputs": [
+          { "data": [1, 2, 3, 6, 7, 8], "dims": [2, 3], "type": "float32" },
+          { "data": [4, 5, 9, 10], "dims": [2, 2], "type": "float32" }
+        ]
+      }
+    ]
+  }
+]

js/web/test/suite-test-list.jsonc

@@ -1409,6 +1409,7 @@
       //"softmax.jsonc",
       "sin.jsonc",
       //"split.jsonc",
+      "split-num-outputs.jsonc",
       "sqrt.jsonc",
       "sub.jsonc",
       "sub_int32.jsonc",

js/web/test/unittests/opset.ts

@@ -8,6 +8,15 @@ import { WEBGL_OP_RESOLVE_RULES } from '../../lib/onnxjs/backends/webgl/op-resol
 import { Graph } from '../../lib/onnxjs/graph';
 import { OpSet, resolveOperator } from '../../lib/onnxjs/opset';
 import { Tensor } from '../../lib/onnxjs/tensor';
+import { DataType } from '../../lib/wasm/wasm-common';
+import { TensorView } from '../../lib/wasm/jsep/tensor-view';
+import {
+  ComputeContext,
+  ComputeContextInputsOutputsMapping,
+  ProgramInfo,
+  TensorInfo,
+} from '../../lib/wasm/jsep/webgpu/types';
+import { parseSplitAttributes, split, SplitAttributes } from '../../lib/wasm/jsep/webgpu/ops/split';
 
 function createTestGraphNode(name: string, opType: string): Graph.Node {
   return { name, opType, inputs: [], outputs: [], attributes: new Attribute(null) };
@@ -106,3 +115,106 @@ describe('#UnitTest# - resolve rules', () => {
     checkConsistency(webglCheckOnlyRules);
   });
 });
+
+describe('#UnitTest# - JSEP Split runtime shapes', () => {
+  const runSplit = (
+    dims: number[],
+    attributes: SplitAttributes,
+    splitSizes?: bigint[] | null,
+  ): readonly TensorInfo[] => {
+    const inputs = [{ dims, dataType: DataType.float }] as unknown as TensorView[];
+    if (splitSizes !== undefined) {
+      inputs.push({
+        dims: splitSizes === null ? [] : [splitSizes.length],
+        dataType: splitSizes === null ? 0 : DataType.int64,
+        getBigInt64Array: () => BigInt64Array.from(splitSizes ?? []),
+      } as unknown as TensorView);
+    }
+    const outputs: TensorInfo[] = [];
+    const context = {
+      inputs,
+      output: (index: number, shape: readonly number[]) => {
+        outputs[index] = { dims: shape, dataType: DataType.float };
+        return 0;
+      },
+      compute: (program: ProgramInfo, mapping: ComputeContextInputsOutputsMapping) => {
+        const programOutputs = program.getRunData(inputs).outputs;
+        programOutputs.forEach((output, index) => {
+          expect(output.dims.every((dim) => dim > 0)).to.equal(true);
+          outputs[mapping.outputs![index]] = output;
+        });
+        return [];
+      },
+    } as unknown as ComputeContext;
+    split(context, attributes);
+    return outputs;
+  };
+
+  const attributes = parseSplitAttributes({ axis: 0, numOutputs: 2, splitSizes: [], isUnevenSplitAllowed: true });
+
+  it('infers even split sizes without graph shapes', () => {
+    expect(runSplit([4], attributes).map((output) => output.dims)).to.deep.equal([[2], [2]]);
+  });
+
+  it('ignores the empty optional split input placeholder', () => {
+    expect(runSplit([4], attributes, null).map((output) => output.dims)).to.deep.equal([[2], [2]]);
+  });
+
+  it('recomputes uneven split sizes for changing input shapes', () => {
+    expect(runSplit([5], attributes).map((output) => output.dims)).to.deep.equal([[3], [2]]);
+    expect(runSplit([7], attributes).map((output) => output.dims)).to.deep.equal([[4], [3]]);
+    expect(attributes.splitSizes).to.deep.equal([]);
+  });
+
+  it('normalizes a negative split axis', () => {
+    const negativeAxis = parseSplitAttributes({ axis: -1, numOutputs: 2, splitSizes: [], isUnevenSplitAllowed: true });
+    expect(runSplit([2, 5], negativeAxis).map((output) => output.dims)).to.deep.equal([
+      [2, 3],
+      [2, 2],
+    ]);
+  });
+
+  it('preserves explicit split sizes from an input tensor', () => {
+    expect(runSplit([5], attributes, [1n, 4n]).map((output) => output.dims)).to.deep.equal([[1], [4]]);
+  });
+
+  it('preserves explicit split sizes from an attribute', () => {
+    const explicitSizes = parseSplitAttributes({ axis: 0, numOutputs: 2, splitSizes: [1, 4] });
+    expect(runSplit([5], explicitSizes).map((output) => output.dims)).to.deep.equal([[1], [4]]);
+  });
+
+  it('requires even splitting without the num_outputs attribute', () => {
+    const evenOnly = parseSplitAttributes({ axis: 0, numOutputs: 2, splitSizes: [] });
+    expect(runSplit([4], evenOnly).map((output) => output.dims)).to.deep.equal([[2], [2]]);
+    expect(() => runSplit([5], evenOnly)).to.throw('evenly divisible');
+  });
+
+  it('allows a zero-sized final output', () => {
+    const sixOutputs = parseSplitAttributes({ axis: 0, numOutputs: 6, splitSizes: [], isUnevenSplitAllowed: true });
+    expect(runSplit([10], sixOutputs).map((output) => output.dims)).to.deep.equal([[2], [2], [2], [2], [2], [0]]);
+    expect(runSplit([10], sixOutputs, null).map((output) => output.dims)).to.deep.equal([[2], [2], [2], [2], [2], [0]]);
+    expect(runSplit([1], attributes).map((output) => output.dims)).to.deep.equal([[1], [0]]);
+  });
+
+  it('allows zero-sized outputs for an empty input', () => {
+    expect(runSplit([0], attributes).map((output) => output.dims)).to.deep.equal([[0], [0]]);
+    expect(runSplit([0, 4], attributes).map((output) => output.dims)).to.deep.equal([
+      [0, 4],
+      [0, 4],
+    ]);
+  });
+
+  it('preserves output positions when explicit split sizes include empty outputs', () => {
+    const explicitSizes = parseSplitAttributes({ axis: 0, numOutputs: 3, splitSizes: [0, 2, 0] });
+    expect(runSplit([2], explicitSizes).map((output) => output.dims)).to.deep.equal([[0], [2], [0]]);
+  });
+
+  it('rejects invalid output counts', () => {
+    const invalidCount = parseSplitAttributes({ axis: 0, numOutputs: 0, splitSizes: [], isUnevenSplitAllowed: true });
+    expect(() => runSplit([4], invalidCount)).to.throw('numOutputs must be positive');
+    const threeOutputs = parseSplitAttributes({ axis: 0, numOutputs: 3, splitSizes: [], isUnevenSplitAllowed: true });
+    expect(() => runSplit([1], threeOutputs)).to.throw('negative split size');
+    const tooManyChunks = parseSplitAttributes({ axis: 0, numOutputs: 7, splitSizes: [], isUnevenSplitAllowed: true });
+    expect(() => runSplit([10], tooManyChunks)).to.throw('negative split size');
+  });
+});

onnxruntime/core/providers/coreml/builders/impl/split_op_builder.cc

@@ -63,7 +63,7 @@ Status SplitOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
     std::unique_ptr<Operation> split_op = model_builder.CreateOperation(node, "split");
     AddOperationInput(*split_op, "axis", model_builder.AddScalarConstant(split_op->type(), "axis", axis));
 
-    if (input_defs.size() > 1) {
+    if (input_defs.size() > 1 && input_defs[1]->Exists()) {
       // if "split" is explicitly provided as an input
       const auto& const_init = *model_builder.GetConstantInitializer(input_defs[1]->Name());
       const Initializer unpacked_tensor(model_builder.GetGraphViewer().GetGraph(), const_init,
@@ -107,7 +107,7 @@ Status SplitOpBuilder::AddToModelBuilderImpl(ModelBuilder& model_builder,
     auto* coreml_splitnd = layer->mutable_splitnd();
     coreml_splitnd->set_axis(axis);
 
-    if (input_defs.size() > 1) {
+    if (input_defs.size() > 1 && input_defs[1]->Exists()) {
       // if "split" is explicitly provided as an input
       // const auto& split_tensor = *model_builder.GetInitializerTensors().at(input_defs[1]->Name());
       const auto& const_init = *model_builder.GetConstantInitializer(input_defs[1]->Name());

onnxruntime/core/providers/cpu/tensor/split.cc

@@ -54,7 +54,6 @@ ONNX_CPU_OPERATOR_VERSIONED_KERNEL(
                                       BuildKernelDefConstraintsFromTypeList<EnabledSplitDataTypes>()),
     Split_1_13);
 
-// TODO: support unequal split and num_outputs
 ONNX_CPU_OPERATOR_KERNEL(
     Split,
     18,