UNPKG

@tensorflow/tfjs-core

Version:

Hardware-accelerated JavaScript library for machine intelligence

42 lines (41 loc) 1.57 kB
/// <amd-module name="@tensorflow/tfjs-core/dist/ops/split" /> import { Tensor } from '../tensor'; import { TensorLike } from '../types'; /** * Splits a `tf.Tensor` into sub tensors. * * If `numOrSizeSplits` is a number, splits `x` along dimension `axis` * into `numOrSizeSplits` smaller tensors. * Requires that `numOrSizeSplits` evenly divides `x.shape[axis]`. * * If `numOrSizeSplits` is a number array, splits `x` into * `numOrSizeSplits.length` pieces. The shape of the `i`-th piece has the * same size as `x` except along dimension `axis` where the size is * `numOrSizeSplits[i]`. * * ```js * const x = tf.tensor2d([1, 2, 3, 4, 5, 6, 7, 8], [2, 4]); * const [a, b] = tf.split(x, 2, 1); * a.print(); * b.print(); * * const [c, d, e] = tf.split(x, [1, 2, 1], 1); * c.print(); * d.print(); * e.print(); * ``` * * @param x The input tensor to split. * @param numOrSizeSplits 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 axis The dimension along which to split. Defaults to 0 (the first * dim). * * @doc {heading: 'Tensors', subheading: 'Slicing and Joining'} */ declare function split_<T extends Tensor>(x: Tensor | TensorLike, numOrSizeSplits: number[] | number, axis?: number): T[]; export declare const split: typeof split_; export {};