From 84b43634b0a33de519db643cf7828b2379350f3a Mon Sep 17 00:00:00 2001 From: Chimuanya Bruce Ibecheozor Date: Mon, 27 Jul 2026 14:59:29 +0100 Subject: [PATCH] Add padSequences utility to tfjs-layers --- tfjs-layers/src/index.ts | 1 + tfjs-layers/src/utils/pad_sequences.ts | 96 +++++++++++++++++++++ tfjs-layers/src/utils/pad_sequences_test.ts | 63 ++++++++++++++ 3 files changed, 160 insertions(+) create mode 100644 tfjs-layers/src/utils/pad_sequences.ts create mode 100644 tfjs-layers/src/utils/pad_sequences_test.ts diff --git a/tfjs-layers/src/index.ts b/tfjs-layers/src/index.ts index fd2e0163131..c8f61141aa5 100644 --- a/tfjs-layers/src/index.ts +++ b/tfjs-layers/src/index.ts @@ -37,3 +37,4 @@ export {ModelAndWeightsConfig, Sequential, SequentialArgs} from './models'; export {LayerVariable} from './variables'; export {version as version_layers} from './version'; export {constraints, initializers, layers, metrics, models, regularizers}; +export {padSequences, PadSequencesOptions} from './utils/pad_sequences'; diff --git a/tfjs-layers/src/utils/pad_sequences.ts b/tfjs-layers/src/utils/pad_sequences.ts new file mode 100644 index 00000000000..e6562596827 --- /dev/null +++ b/tfjs-layers/src/utils/pad_sequences.ts @@ -0,0 +1,96 @@ +/** + * @license + * Copyright 2026 Google LLC + * + * Use of this source code is governed by an MIT-style + * license that can be found in the LICENSE file or at + * https://opensource.org/licenses/MIT. + * ============================================================================= + */ + +import {Tensor, tensor2d} from '@tensorflow/tfjs-core'; + +/** + * Configuration options for `padSequences`. + */ +export interface PadSequencesOptions { + /** + * Maximum length of all sequences. If omitted, it defaults to the + * length of the longest sequence in the input. + */ + maxlen?: number; + /** + * String: 'pre' or 'post'. Pad either before or after each sequence. + * Defaults to 'pre'. + */ + padding?: 'pre'|'post'; + /** + * String: 'pre' or 'post'. Remove values from sequences larger than + * `maxlen`, either at the beginning or at the end of the sequences. + * Defaults to 'pre'. + */ + truncating?: 'pre'|'post'; + /** + * Float or String: Pad value (must match output dtype). + * Defaults to 0. + */ + value?: number; + /** + * Data type of the output Tensor. + * Defaults to 'int32'. + */ + dtype?: 'int32'|'float32'; +} + +/** + * Pads or truncates a list of sequences (nested numeric arrays) to the + * same length. + * + * This function mimics the behavior of Python Keras + * `tf.keras.utils.pad_sequences`. + * + * @param sequences A 2D array of numeric sequences (e.g., tokenized + * integer arrays). + * @param options Configuration options for padding, truncating, max + * length, and data type. + * @returns A 2D Tensor (`Tensor2D`) containing the padded/truncated + * sequences with shape `[sequences.length, maxlen]`. + */ +export function padSequences( + sequences: number[][], options: PadSequencesOptions = {}): Tensor { + // Extract configuration options with sensible Keras-compatible defaults + const {padding = 'pre', truncating = 'pre', value = 0, dtype = 'int32'} = + options; + + // Calculate maximum length: use user-defined maxlen if provided, + // otherwise find the length of the longest input sequence. + const lengths = sequences.map(s => s.length); + const maxlen = options.maxlen ?? Math.max(...lengths, 0); + + // Process each sequence individually + const padded = sequences.map(seq => { + let current = [...seq]; + + // Truncate sequence if it exceeds maxlen + if (current.length > maxlen) { + current = truncating === 'post' ? + current.slice(0, maxlen) : // Keep first `maxlen` + current.slice(current.length - maxlen); // Keep last `maxlen` + } + + // Pad sequence if it is shorter than maxlen + if (current.length < maxlen) { + const padCount = maxlen - current.length; + const padArray = new Array(padCount).fill(value); + + current = padding === 'post' ? + [...current, ...padArray] : // Append pad values to end + [...padArray, ...current]; // Prepend pad values to start + } + + return current; + }); + + // Convert the 2D padded JavaScript array into a 2D TensorFlow Tensor + return tensor2d(padded, [sequences.length, maxlen], dtype); +} diff --git a/tfjs-layers/src/utils/pad_sequences_test.ts b/tfjs-layers/src/utils/pad_sequences_test.ts new file mode 100644 index 00000000000..84111f1ea7d --- /dev/null +++ b/tfjs-layers/src/utils/pad_sequences_test.ts @@ -0,0 +1,63 @@ +/** + * @license + * Copyright 2026 Google LLC + * + * Use of this source code is governed by an MIT-style + * license that can be found in the LICENSE file or at + * https://opensource.org/licenses/MIT. + * ============================================================================= + */ + +import * as tfc from '@tensorflow/tfjs-core'; + +import {padSequences} from './pad_sequences'; +import {expectTensorsClose} from './test_utils'; + +describe('padSequences', () => { + it('default pre-padding and automatic maxlen calculation', () => { + const seqs = [[1, 2], [3, 4, 5]]; + const y = padSequences(seqs); + + expect(y.shape).toEqual([2, 3]); + expectTensorsClose( + y, tfc.tensor2d([[0, 1, 2], [3, 4, 5]], [2, 3], 'int32')); + }); + + it('custom post-padding', () => { + const seqs = [[1, 2], [3, 4, 5]]; + const y = padSequences(seqs, {maxlen: 3, padding: 'post'}); + + expectTensorsClose( + y, tfc.tensor2d([[1, 2, 0], [3, 4, 5]], [2, 3], 'int32')); + }); + + it('pre-truncating when sequence length exceeds maxlen', () => { + const seqs = [[1, 2, 3, 4, 5]]; + const y = padSequences(seqs, {maxlen: 3, truncating: 'pre'}); + + expectTensorsClose(y, tfc.tensor2d([[3, 4, 5]], [1, 3], 'int32')); + }); + + it('post-truncating when sequence length exceeds maxlen', () => { + const seqs = [[1, 2, 3, 4, 5]]; + const y = padSequences(seqs, {maxlen: 3, truncating: 'post'}); + + expectTensorsClose(y, tfc.tensor2d([[1, 2, 3]], [1, 3], 'int32')); + }); + + it('custom pad value and float32 dtype', () => { + const seqs = [[1.5, 2.5]]; + const y = padSequences(seqs, {maxlen: 4, value: -1, dtype: 'float32'}); + + expect(y.dtype).toEqual('float32'); + expectTensorsClose( + y, tfc.tensor2d([[-1, -1, 1.5, 2.5]], [1, 4], 'float32')); + }); + + it('handles empty input sequences', () => { + const seqs: number[][] = []; + const y = padSequences(seqs, {maxlen: 5}); + + expect(y.shape).toEqual([0, 5]); + }); +});