import { Order } from 'blockly/javascript';
/**
 * @module SlicingAndJoiningTensors
 * @category AI
 * @subcategory Tensor
 */
/**
 * Splits a tf.Tensor into sub tensors.
 *
 * @param {Tensor} tensor The input tensor to split.
 * @param {Array<number> | Number} size Either an integer indicating the number of splits along the axis or an array of integers containing the sizes of each output tensor along the axis. If a number then it must evenly divide x.shape[axis]; otherwise the sum of sizes must match x.shape[axis]. Can contain one -1 indicating that dimension is to be inferred.
 * @param {Number} axis The dimension along which to split. Defaults to 0 (the first dim).
 * @returns {Tensor} tf.split(tensor, size, axis).
 */
function splitTensor(block, generator) {
    const tensor = generator.valueToCode(block, 'TENSOR_VALUES', Order.NONE) || ''; //tensor to split
    const size = generator.valueToCode(block, 'SIZE', Order.NONE) || ''; //numOrSizeSplits
    const axis = generator.valueToCode(block, 'AXIS', Order.NONE) || ''; //axis
    return [
        `tf.split(${tensor}, ${size} ${parseInt(axis) === 0 || parseInt(axis) === 1 ? `, ${axis}` : ''})`,
        Order.VOID,
    ];
}
export const slicingAndJoiningBlockGenerator = {
    ['split']: splitTensor,
};
//# sourceMappingURL=generators.js.map