diff --git a/packages/tasks/src/tasks/feature-extraction/inference.ts b/packages/tasks/src/tasks/feature-extraction/inference.ts
index 9317d38ac3..9049a45c46 100644
--- a/packages/tasks/src/tasks/feature-extraction/inference.ts
+++ b/packages/tasks/src/tasks/feature-extraction/inference.ts
@@ -12,6 +12,11 @@ export type FeatureExtractionOutput = Array<number[]>;
  * https://github.com/huggingface/huggingface.js/blob/main/packages/tasks/scripts/inference-tei-import.ts.
  */
 export interface FeatureExtractionInput {
+	/**
+	 * The number of dimensions that the output embeddings should have. If not set, the original
+	 * shape of the representation will be returned instead.
+	 */
+	dimensions?: number;
 	/**
 	 * The text or list of texts to embed.
 	 */
diff --git a/packages/tasks/src/tasks/feature-extraction/spec/input.json b/packages/tasks/src/tasks/feature-extraction/spec/input.json
index 55f58e5b88..5cd311946b 100644
--- a/packages/tasks/src/tasks/feature-extraction/spec/input.json
+++ b/packages/tasks/src/tasks/feature-extraction/spec/input.json
@@ -6,6 +6,14 @@
 	"type": "object",
 	"required": ["inputs"],
 	"properties": {
+		"dimensions": {
+			"type": "integer",
+			"description": "The number of dimensions that the output embeddings should have. If not set, the original\nshape of the representation will be returned instead.",
+			"default": "null",
+			"example": "null",
+			"nullable": true,
+			"minimum": 0
+		},
 		"inputs": {
 			"title": "FeatureExtractionInputs",
 			"description": "The text or list of texts to embed.",