UNPKG

1.51 kBTypeScriptView Raw
1import { Tensor } from '../tensor';
2import { TensorLike } from '../types';
3/**
4 * Splits a `tf.Tensor` into sub tensors.
5 *
6 * If `numOrSizeSplits` is a number, splits `x` along dimension `axis`
7 * into `numOrSizeSplits` smaller tensors.
8 * Requires that `numOrSizeSplits` evenly divides `x.shape[axis]`.
9 *
10 * If `numOrSizeSplits` is a number array, splits `x` into
11 * `numOrSizeSplits.length` pieces. The shape of the `i`-th piece has the
12 * same size as `x` except along dimension `axis` where the size is
13 * `numOrSizeSplits[i]`.
14 *
15 * ```js
16 * const x = tf.tensor2d([1, 2, 3, 4, 5, 6, 7, 8], [2, 4]);
17 * const [a, b] = tf.split(x, 2, 1);
18 * a.print();
19 * b.print();
20 *
21 * const [c, d, e] = tf.split(x, [1, 2, 1], 1);
22 * c.print();
23 * d.print();
24 * e.print();
25 * ```
26 *
27 * @param x The input tensor to split.
28 * @param numOrSizeSplits Either an integer indicating the number of
29 * splits along the axis or an array of integers containing the sizes of
30 * each output tensor along the axis. If a number then it must evenly divide
31 * `x.shape[axis]`; otherwise the sum of sizes must match `x.shape[axis]`.
32 * Can contain one -1 indicating that dimension is to be inferred.
33 * @param axis The dimension along which to split. Defaults to 0 (the first
34 * dim).
35 *
36 * @doc {heading: 'Tensors', subheading: 'Slicing and Joining'}
37 */
38declare function split_<T extends Tensor>(x: Tensor | TensorLike, numOrSizeSplits: number[] | number, axis?: number): T[];
39export declare const split: typeof split_;
40export {};