From a2cf0b41dd009b82555cafa04af50ee1bf0c6013 Mon Sep 17 00:00:00 2001
From: Linchenn <linchan@google.com>
Date: Tue, 30 Aug 2022 18:00:44 -0700
Subject: [PATCH 01/10] Add WEBGL_MAX_TEXTURE_DIMENSION_FOR_1D_TEXTURE

---
 tfjs-backend-webgl/src/flags_webgl.ts      | 10 ++++++++
 tfjs-backend-webgl/src/flags_webgl_test.ts | 16 ++++++++++++
 tfjs-backend-webgl/src/webgl_util.ts       | 30 ++++++++++++++++------
 3 files changed, 48 insertions(+), 8 deletions(-)

diff --git a/tfjs-backend-webgl/src/flags_webgl.ts b/tfjs-backend-webgl/src/flags_webgl.ts
index dbfbcf95c72..3a1651dc71e 100644
--- a/tfjs-backend-webgl/src/flags_webgl.ts
+++ b/tfjs-backend-webgl/src/flags_webgl.ts
@@ -246,3 +246,13 @@ ENV.registerFlag('WEBGL_EXP_CONV', () => false);
  * software WebGL will be used.
  */
 ENV.registerFlag('SOFTWARE_WEBGL_ENABLED', () => ENV.getBool('IS_TEST'));
+
+/**
+ * For 1D texture (physical height or physical width is 1), if the length of any
+ * texture edges exceed the threshold, the other edge will be increased to 2.
+ *
+ * This flag is used to help some GPUs that could not provide correct
+ * interpolations for long skinny triangles.
+ * https://github.com/tensorflow/tfjs/issues/6775
+ */
+ENV.registerFlag('WEBGL_MAX_TEXTURE_DIMENSION_FOR_1D_TEXTURE', () => Infinity);
diff --git a/tfjs-backend-webgl/src/flags_webgl_test.ts b/tfjs-backend-webgl/src/flags_webgl_test.ts
index 3ce0dc4b0ca..0c2d2e84de2 100644
--- a/tfjs-backend-webgl/src/flags_webgl_test.ts
+++ b/tfjs-backend-webgl/src/flags_webgl_test.ts
@@ -444,3 +444,19 @@ describe('WEBGL_EXP_CONV', () => {
     expect(tf.env().getBool('WEBGL_EXP_CONV')).toBe(false);
   });
 });
+
+const MAX_1D_TEX_DIM_FLAG = 'WEBGL_MAX_TEXTURE_DIMENSION_FOR_1D_TEXTURE';
+
+describeWithFlags(MAX_1D_TEX_DIM_FLAG, WEBGL_ENVS, () => {
+  beforeEach(() => tf.env().reset());
+  afterAll(() => tf.env().reset());
+
+  it(`returns correct value when ${MAX_1D_TEX_DIM_FLAG} is set`, () => {
+    tf.env().set(MAX_1D_TEX_DIM_FLAG, 2048);
+    expect(tf.env().getNumber(MAX_1D_TEX_DIM_FLAG)).toBe(2048);
+  });
+
+  it(`returns default when ${MAX_1D_TEX_DIM_FLAG} is not set`, () => {
+    expect(tf.env().getNumber(MAX_1D_TEX_DIM_FLAG)).toBe(Infinity);
+  });
+});
diff --git a/tfjs-backend-webgl/src/webgl_util.ts b/tfjs-backend-webgl/src/webgl_util.ts
index 52c8bd4d8df..3df8dc2041a 100644
--- a/tfjs-backend-webgl/src/webgl_util.ts
+++ b/tfjs-backend-webgl/src/webgl_util.ts
@@ -395,29 +395,30 @@ export function getTextureShapeFromLogicalShape(
   }
 
   let size = util.sizeFromShape(logShape);
+  let textureShape: [number, number];
   if (logShape.length <= 1 && size <= maxTexSize) {
-    return [1, size];
+    textureShape = [1, size];
   } else if (
       logShape.length === 2 && logShape[0] <= maxTexSize &&
       logShape[1] <= maxTexSize) {
-    return logShape as [number, number];
+    textureShape = logShape as [number, number];
   } else if (
       logShape.length === 3 && logShape[0] * logShape[1] <= maxTexSize &&
       logShape[2] <= maxTexSize) {
-    return [logShape[0] * logShape[1], logShape[2]];
+    textureShape = [logShape[0] * logShape[1], logShape[2]];
   } else if (
       logShape.length === 3 && logShape[0] <= maxTexSize &&
       logShape[1] * logShape[2] <= maxTexSize) {
-    return [logShape[0], logShape[1] * logShape[2]];
+    textureShape = [logShape[0], logShape[1] * logShape[2]];
   } else if (
       logShape.length === 4 &&
       logShape[0] * logShape[1] * logShape[2] <= maxTexSize &&
       logShape[3] <= maxTexSize) {
-    return [logShape[0] * logShape[1] * logShape[2], logShape[3]];
+    textureShape = [logShape[0] * logShape[1] * logShape[2], logShape[3]];
   } else if (
       logShape.length === 4 && logShape[0] <= maxTexSize &&
       logShape[1] * logShape[2] * logShape[3] <= maxTexSize) {
-    return [logShape[0], logShape[1] * logShape[2] * logShape[3]];
+    textureShape = [logShape[0], logShape[1] * logShape[2] * logShape[3]];
   } else {
     if (isPacked) {
       // For packed textures size equals the number of channels required to
@@ -432,10 +433,23 @@ export function getTextureShapeFromLogicalShape(
         [rows, cols] = getRowsCols(logShape);
       }
       size = batchDim * (rows / 2) * (cols / 2);
-      return util.sizeToSquarishShape(size).map(d => d * 2) as [number, number];
+      textureShape =
+          util.sizeToSquarishShape(size).map(d => d * 2) as [number, number];
+    } else {
+      textureShape = util.sizeToSquarishShape(size);
     }
-    return util.sizeToSquarishShape(size);
   }
+
+  const maxDimFor1DTex =
+      env().getNumber('WEBGL_MAX_TEXTURE_DIMENSION_FOR_1D_TEXTURE') *
+      (isPacked ? 2 : 1);
+  if (textureShape[0] > maxDimFor1DTex || textureShape[1] > maxDimFor1DTex) {
+    // For 1D texture, if the length exceeds maxDimFor1DTex, the texture will be
+    // upgraded to 2D with a physical shape as [length, 2] or [2, length].
+    textureShape[0] = Math.max(isPacked ? 4 : 2, textureShape[0]);
+    textureShape[1] = Math.max(isPacked ? 4 : 2, textureShape[1]);
+  }
+  return textureShape;
 }
 
 function isEven(n: number): boolean {

From 0c3ef563a0ac6c8298fd627cb1d1f9357fd35b29 Mon Sep 17 00:00:00 2001
From: Linchenn <linchan@google.com>
Date: Wed, 31 Aug 2022 10:55:46 -0700
Subject: [PATCH 02/10] update flag

---
 tfjs-backend-webgl/src/flags_webgl.ts      | 22 ++++++--
 tfjs-backend-webgl/src/flags_webgl_test.ts | 32 ++++++++---
 tfjs-backend-webgl/src/webgl_util.ts       | 65 ++++++++++++----------
 3 files changed, 79 insertions(+), 40 deletions(-)

diff --git a/tfjs-backend-webgl/src/flags_webgl.ts b/tfjs-backend-webgl/src/flags_webgl.ts
index 3a1651dc71e..264b55fd439 100644
--- a/tfjs-backend-webgl/src/flags_webgl.ts
+++ b/tfjs-backend-webgl/src/flags_webgl.ts
@@ -248,11 +248,23 @@ ENV.registerFlag('WEBGL_EXP_CONV', () => false);
 ENV.registerFlag('SOFTWARE_WEBGL_ENABLED', () => ENV.getBool('IS_TEST'));
 
 /**
- * For 1D texture (physical height or physical width is 1), if the length of any
- * texture edges exceed the threshold, the other edge will be increased to 2.
+ * For narrow texture (physical height or physical width is 1), if the length of
+ * any texture edges exceed the threshold, the texture will be reshaped to be
+ * more squarish.
  *
  * This flag is used to help some GPUs that could not provide correct
- * interpolations for long skinny triangles.
- * https://github.com/tensorflow/tfjs/issues/6775
+ * interpolations for long skinny triangles. We found Mali GPU probably has this
+ * problem: https://github.com/tensorflow/tfjs/issues/6775.
  */
-ENV.registerFlag('WEBGL_MAX_TEXTURE_DIMENSION_FOR_1D_TEXTURE', () => Infinity);
+ENV.registerFlag('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE', () => Infinity);
+
+/**
+ * If the flag is set to true, the max size of the narrow texture will be auto
+ * computed and it will be considerred as a threshold to reshape the narrow
+ * texture to be more squarish.
+ *
+ * This flag is used to help some GPUs that could not provide correct
+ * interpolations for long skinny triangles. We found Mali GPU probably has this
+ * problem: https://github.com/tensorflow/tfjs/issues/6775.
+ */
+ENV.registerFlag('WEBGL_AUTO_RESHAPE_NARROW_TEXTURE_FOR_MALI_GPU', () => false);
diff --git a/tfjs-backend-webgl/src/flags_webgl_test.ts b/tfjs-backend-webgl/src/flags_webgl_test.ts
index 0c2d2e84de2..62d6cf354fd 100644
--- a/tfjs-backend-webgl/src/flags_webgl_test.ts
+++ b/tfjs-backend-webgl/src/flags_webgl_test.ts
@@ -445,18 +445,36 @@ describe('WEBGL_EXP_CONV', () => {
   });
 });
 
-const MAX_1D_TEX_DIM_FLAG = 'WEBGL_MAX_TEXTURE_DIMENSION_FOR_1D_TEXTURE';
+const MAX_SIZE_FOR_NARROR_TEX_FLAG = 'WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE';
 
-describeWithFlags(MAX_1D_TEX_DIM_FLAG, WEBGL_ENVS, () => {
+describeWithFlags(MAX_SIZE_FOR_NARROR_TEX_FLAG, WEBGL_ENVS, () => {
   beforeEach(() => tf.env().reset());
   afterAll(() => tf.env().reset());
 
-  it(`returns correct value when ${MAX_1D_TEX_DIM_FLAG} is set`, () => {
-    tf.env().set(MAX_1D_TEX_DIM_FLAG, 2048);
-    expect(tf.env().getNumber(MAX_1D_TEX_DIM_FLAG)).toBe(2048);
+  it(`returns correct value when ${MAX_SIZE_FOR_NARROR_TEX_FLAG} is set`,
+     () => {
+       tf.env().set(MAX_SIZE_FOR_NARROR_TEX_FLAG, 2048);
+       expect(tf.env().getNumber(MAX_SIZE_FOR_NARROR_TEX_FLAG)).toBe(2048);
+     });
+
+  it(`returns default when ${MAX_SIZE_FOR_NARROR_TEX_FLAG} is not set`, () => {
+    expect(tf.env().getNumber(MAX_SIZE_FOR_NARROR_TEX_FLAG)).toBe(Infinity);
   });
+});
+
+const AUTO_RESHAPE_NARROW_TEX_FLAG = 'WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE';
+
+describeWithFlags(AUTO_RESHAPE_NARROW_TEX_FLAG, WEBGL_ENVS, () => {
+  beforeEach(() => tf.env().reset());
+  afterAll(() => tf.env().reset());
+
+  it(`returns correct value when ${AUTO_RESHAPE_NARROW_TEX_FLAG} is set`,
+     () => {
+       tf.env().set(AUTO_RESHAPE_NARROW_TEX_FLAG, true);
+       expect(tf.env().getBool(AUTO_RESHAPE_NARROW_TEX_FLAG)).toBe(true);
+     });
 
-  it(`returns default when ${MAX_1D_TEX_DIM_FLAG} is not set`, () => {
-    expect(tf.env().getNumber(MAX_1D_TEX_DIM_FLAG)).toBe(Infinity);
+  it(`returns default when ${AUTO_RESHAPE_NARROW_TEX_FLAG} is not set`, () => {
+    expect(tf.env().getBool(AUTO_RESHAPE_NARROW_TEX_FLAG)).toBe(false);
   });
 });
diff --git a/tfjs-backend-webgl/src/webgl_util.ts b/tfjs-backend-webgl/src/webgl_util.ts
index 3df8dc2041a..82ef97d04ac 100644
--- a/tfjs-backend-webgl/src/webgl_util.ts
+++ b/tfjs-backend-webgl/src/webgl_util.ts
@@ -368,8 +368,17 @@ export function getShapeAs3D(shape: number[]): [number, number, number] {
 export function getTextureShapeFromLogicalShape(
     logShape: number[], isPacked = false): [number, number] {
   let maxTexSize = env().getNumber('WEBGL_MAX_TEXTURE_SIZE');
+  let maxSizeForNarrorTex =
+      env().getNumber('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE');
+
+  if (maxSizeForNarrorTex === Infinity &&
+      env().getBool('WEBGL_AUTO_RESHAPE_NARROW_TEXTURE_FOR_MALI_GPU')) {
+    maxSizeForNarrorTex = maxTexSize / 2;
+  }
+
   if (isPacked) {
     maxTexSize = maxTexSize * 2;
+    maxSizeForNarrorTex = maxSizeForNarrorTex * 2;
 
     // This logic ensures we accurately count the number of packed texels needed
     // to accommodate the tensor. We can only pack values in the same texel if
@@ -388,6 +397,12 @@ export function getTextureShapeFromLogicalShape(
     }
   }
 
+  const isLongNarrowTex =
+      (textureShape: [number, number]) => {
+        return Math.max(...textureShape) > maxSizeForNarrorTex &&
+            Math.min(...textureShape) <= (isPacked ? 2 : 1);
+      }
+
   // If logical shape is 2, we don't squeeze, since we want to match physical.
   if (logShape.length !== 2) {
     const squeezeResult = util.squeezeShape(logShape);
@@ -395,30 +410,37 @@ export function getTextureShapeFromLogicalShape(
   }
 
   let size = util.sizeFromShape(logShape);
-  let textureShape: [number, number];
-  if (logShape.length <= 1 && size <= maxTexSize) {
-    textureShape = [1, size];
+  if (logShape.length <= 1 && size <= maxTexSize &&
+      size <= maxSizeForNarrorTex) {
+    return [1, size];
   } else if (
       logShape.length === 2 && logShape[0] <= maxTexSize &&
-      logShape[1] <= maxTexSize) {
-    textureShape = logShape as [number, number];
+      logShape[1] <= maxTexSize &&
+      !isLongNarrowTex(logShape as [number, number])) {
+    return logShape as [number, number];
   } else if (
       logShape.length === 3 && logShape[0] * logShape[1] <= maxTexSize &&
-      logShape[2] <= maxTexSize) {
-    textureShape = [logShape[0] * logShape[1], logShape[2]];
+      logShape[2] <= maxTexSize &&
+      !isLongNarrowTex([logShape[0] * logShape[1], logShape[2]])) {
+    return [logShape[0] * logShape[1], logShape[2]];
   } else if (
       logShape.length === 3 && logShape[0] <= maxTexSize &&
-      logShape[1] * logShape[2] <= maxTexSize) {
-    textureShape = [logShape[0], logShape[1] * logShape[2]];
+      logShape[1] * logShape[2] <= maxTexSize &&
+      !isLongNarrowTex([logShape[0], logShape[1] * logShape[2]])) {
+    return [logShape[0], logShape[1] * logShape[2]];
   } else if (
       logShape.length === 4 &&
       logShape[0] * logShape[1] * logShape[2] <= maxTexSize &&
-      logShape[3] <= maxTexSize) {
-    textureShape = [logShape[0] * logShape[1] * logShape[2], logShape[3]];
+      logShape[3] <= maxTexSize &&
+      !isLongNarrowTex(
+          [logShape[0] * logShape[1] * logShape[2], logShape[3]])) {
+    return [logShape[0] * logShape[1] * logShape[2], logShape[3]];
   } else if (
       logShape.length === 4 && logShape[0] <= maxTexSize &&
-      logShape[1] * logShape[2] * logShape[3] <= maxTexSize) {
-    textureShape = [logShape[0], logShape[1] * logShape[2] * logShape[3]];
+      logShape[1] * logShape[2] * logShape[3] <= maxTexSize &&
+      !isLongNarrowTex(
+          [logShape[0], logShape[1] * logShape[2] * logShape[3]])) {
+    return [logShape[0], logShape[1] * logShape[2] * logShape[3]];
   } else {
     if (isPacked) {
       // For packed textures size equals the number of channels required to
@@ -433,23 +455,10 @@ export function getTextureShapeFromLogicalShape(
         [rows, cols] = getRowsCols(logShape);
       }
       size = batchDim * (rows / 2) * (cols / 2);
-      textureShape =
-          util.sizeToSquarishShape(size).map(d => d * 2) as [number, number];
-    } else {
-      textureShape = util.sizeToSquarishShape(size);
+      return util.sizeToSquarishShape(size).map(d => d * 2) as [number, number];
     }
+    return util.sizeToSquarishShape(size);
   }
-
-  const maxDimFor1DTex =
-      env().getNumber('WEBGL_MAX_TEXTURE_DIMENSION_FOR_1D_TEXTURE') *
-      (isPacked ? 2 : 1);
-  if (textureShape[0] > maxDimFor1DTex || textureShape[1] > maxDimFor1DTex) {
-    // For 1D texture, if the length exceeds maxDimFor1DTex, the texture will be
-    // upgraded to 2D with a physical shape as [length, 2] or [2, length].
-    textureShape[0] = Math.max(isPacked ? 4 : 2, textureShape[0]);
-    textureShape[1] = Math.max(isPacked ? 4 : 2, textureShape[1]);
-  }
-  return textureShape;
 }
 
 function isEven(n: number): boolean {

From 16770325837b02b62374f6eea87ec9dddfa9b71d Mon Sep 17 00:00:00 2001
From: Linchenn <linchan@google.com>
Date: Wed, 31 Aug 2022 11:27:50 -0700
Subject: [PATCH 03/10] add tests

---
 tfjs-backend-webgl/src/flags_webgl.ts      |  3 +-
 tfjs-backend-webgl/src/flags_webgl_test.ts | 15 +++++-----
 tfjs-backend-webgl/src/webgl_util.ts       |  5 ++--
 tfjs-backend-webgl/src/webgl_util_test.ts  | 35 ++++++++++++++++++++++
 4 files changed, 48 insertions(+), 10 deletions(-)

diff --git a/tfjs-backend-webgl/src/flags_webgl.ts b/tfjs-backend-webgl/src/flags_webgl.ts
index 264b55fd439..427118e22a4 100644
--- a/tfjs-backend-webgl/src/flags_webgl.ts
+++ b/tfjs-backend-webgl/src/flags_webgl.ts
@@ -267,4 +267,5 @@ ENV.registerFlag('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE', () => Infinity);
  * interpolations for long skinny triangles. We found Mali GPU probably has this
  * problem: https://github.com/tensorflow/tfjs/issues/6775.
  */
-ENV.registerFlag('WEBGL_AUTO_RESHAPE_NARROW_TEXTURE_FOR_MALI_GPU', () => false);
+ENV.registerFlag(
+    'WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU', () => false);
diff --git a/tfjs-backend-webgl/src/flags_webgl_test.ts b/tfjs-backend-webgl/src/flags_webgl_test.ts
index 62d6cf354fd..87432992493 100644
--- a/tfjs-backend-webgl/src/flags_webgl_test.ts
+++ b/tfjs-backend-webgl/src/flags_webgl_test.ts
@@ -462,19 +462,20 @@ describeWithFlags(MAX_SIZE_FOR_NARROR_TEX_FLAG, WEBGL_ENVS, () => {
   });
 });
 
-const AUTO_RESHAPE_NARROW_TEX_FLAG = 'WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE';
+const AUTO_SQUARIFY_NARROW_TEX_FLAG =
+    'WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU';
 
-describeWithFlags(AUTO_RESHAPE_NARROW_TEX_FLAG, WEBGL_ENVS, () => {
+describeWithFlags(AUTO_SQUARIFY_NARROW_TEX_FLAG, WEBGL_ENVS, () => {
   beforeEach(() => tf.env().reset());
   afterAll(() => tf.env().reset());
 
-  it(`returns correct value when ${AUTO_RESHAPE_NARROW_TEX_FLAG} is set`,
+  it(`returns correct value when ${AUTO_SQUARIFY_NARROW_TEX_FLAG} is set`,
      () => {
-       tf.env().set(AUTO_RESHAPE_NARROW_TEX_FLAG, true);
-       expect(tf.env().getBool(AUTO_RESHAPE_NARROW_TEX_FLAG)).toBe(true);
+       tf.env().set(AUTO_SQUARIFY_NARROW_TEX_FLAG, true);
+       expect(tf.env().getBool(AUTO_SQUARIFY_NARROW_TEX_FLAG)).toBe(true);
      });
 
-  it(`returns default when ${AUTO_RESHAPE_NARROW_TEX_FLAG} is not set`, () => {
-    expect(tf.env().getBool(AUTO_RESHAPE_NARROW_TEX_FLAG)).toBe(false);
+  it(`returns default when ${AUTO_SQUARIFY_NARROW_TEX_FLAG} is not set`, () => {
+    expect(tf.env().getBool(AUTO_SQUARIFY_NARROW_TEX_FLAG)).toBe(false);
   });
 });
diff --git a/tfjs-backend-webgl/src/webgl_util.ts b/tfjs-backend-webgl/src/webgl_util.ts
index 82ef97d04ac..85b57dc853d 100644
--- a/tfjs-backend-webgl/src/webgl_util.ts
+++ b/tfjs-backend-webgl/src/webgl_util.ts
@@ -370,9 +370,8 @@ export function getTextureShapeFromLogicalShape(
   let maxTexSize = env().getNumber('WEBGL_MAX_TEXTURE_SIZE');
   let maxSizeForNarrorTex =
       env().getNumber('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE');
-
   if (maxSizeForNarrorTex === Infinity &&
-      env().getBool('WEBGL_AUTO_RESHAPE_NARROW_TEXTURE_FOR_MALI_GPU')) {
+      env().getBool('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU')) {
     maxSizeForNarrorTex = maxTexSize / 2;
   }
 
@@ -397,6 +396,8 @@ export function getTextureShapeFromLogicalShape(
     }
   }
 
+  // returns true if one edge length is 1 (or 2, if packed), while another edge
+  // length exceeds maxSizeForNarrorTex.
   const isLongNarrowTex =
       (textureShape: [number, number]) => {
         return Math.max(...textureShape) > maxSizeForNarrorTex &&
diff --git a/tfjs-backend-webgl/src/webgl_util_test.ts b/tfjs-backend-webgl/src/webgl_util_test.ts
index 1377a4b25f0..17c666d2dc5 100644
--- a/tfjs-backend-webgl/src/webgl_util_test.ts
+++ b/tfjs-backend-webgl/src/webgl_util_test.ts
@@ -114,6 +114,41 @@ describeWithFlags('getTextureShapeFromLogicalShape packed', WEBGL_ENVS, () => {
     tf.env().set('WEBGL_MAX_TEXTURE_SIZE', max);
     expect(texShape).toEqual([6, 4]);
   });
+
+  it('squarified long narrow texture shapes', () => {
+    const isPacked = true;
+    const max = tf.env().getNumber('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE');
+
+    tf.env().set('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE', 5);
+    const logicalShape = [1, 16];
+    const texShape =
+        webgl_util.getTextureShapeFromLogicalShape(logicalShape, isPacked);
+
+    tf.env().set('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE', max);
+    expect(texShape).toEqual([6, 6]);
+  });
+
+  it('auto squarified long narrow texture shapes', () => {
+    const isPacked = true;
+    const max = tf.env().getNumber('WEBGL_MAX_TEXTURE_SIZE');
+    const maxForNarrowTex =
+        tf.env().getNumber('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE');
+    const autoSquarify = tf.env().getNumber(
+        'WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU');
+
+    tf.env().set('WEBGL_MAX_TEXTURE_SIZE', 6);
+    tf.env().set('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU', true);
+    tf.env().set('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE', Infinity);
+    const logicalShape = [1, 16];
+    const texShape =
+        webgl_util.getTextureShapeFromLogicalShape(logicalShape, isPacked);
+
+    tf.env().set('WEBGL_MAX_TEXTURE_SIZE', max);
+    tf.env().set('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE', maxForNarrowTex);
+    tf.env().set(
+        'WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU', autoSquarify);
+    expect(texShape).toEqual([6, 6]);
+  });
 });
 
 describeWithFlags('isReshapeFree', WEBGL_ENVS, () => {

From 45dce514c2dd6199075321a38065983ee8926e1e Mon Sep 17 00:00:00 2001
From: Linchenn <linchan@google.com>
Date: Wed, 31 Aug 2022 14:43:39 -0700
Subject: [PATCH 04/10] rename flag

---
 tfjs-backend-webgl/src/flags_webgl.ts      | 3 +--
 tfjs-backend-webgl/src/flags_webgl_test.ts | 2 +-
 tfjs-backend-webgl/src/webgl_util.ts       | 5 +++--
 tfjs-backend-webgl/src/webgl_util_test.ts  | 9 ++++-----
 4 files changed, 9 insertions(+), 10 deletions(-)

diff --git a/tfjs-backend-webgl/src/flags_webgl.ts b/tfjs-backend-webgl/src/flags_webgl.ts
index 427118e22a4..6846e1d64a8 100644
--- a/tfjs-backend-webgl/src/flags_webgl.ts
+++ b/tfjs-backend-webgl/src/flags_webgl.ts
@@ -267,5 +267,4 @@ ENV.registerFlag('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE', () => Infinity);
  * interpolations for long skinny triangles. We found Mali GPU probably has this
  * problem: https://github.com/tensorflow/tfjs/issues/6775.
  */
-ENV.registerFlag(
-    'WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU', () => false);
+ENV.registerFlag('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE', () => false);
diff --git a/tfjs-backend-webgl/src/flags_webgl_test.ts b/tfjs-backend-webgl/src/flags_webgl_test.ts
index 87432992493..a8568863efc 100644
--- a/tfjs-backend-webgl/src/flags_webgl_test.ts
+++ b/tfjs-backend-webgl/src/flags_webgl_test.ts
@@ -463,7 +463,7 @@ describeWithFlags(MAX_SIZE_FOR_NARROR_TEX_FLAG, WEBGL_ENVS, () => {
 });
 
 const AUTO_SQUARIFY_NARROW_TEX_FLAG =
-    'WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU';
+    'WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE';
 
 describeWithFlags(AUTO_SQUARIFY_NARROW_TEX_FLAG, WEBGL_ENVS, () => {
   beforeEach(() => tf.env().reset());
diff --git a/tfjs-backend-webgl/src/webgl_util.ts b/tfjs-backend-webgl/src/webgl_util.ts
index 85b57dc853d..85c5bb62726 100644
--- a/tfjs-backend-webgl/src/webgl_util.ts
+++ b/tfjs-backend-webgl/src/webgl_util.ts
@@ -371,7 +371,7 @@ export function getTextureShapeFromLogicalShape(
   let maxSizeForNarrorTex =
       env().getNumber('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE');
   if (maxSizeForNarrorTex === Infinity &&
-      env().getBool('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU')) {
+      env().getBool('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE')) {
     maxSizeForNarrorTex = maxTexSize / 2;
   }
 
@@ -401,7 +401,8 @@ export function getTextureShapeFromLogicalShape(
   const isLongNarrowTex =
       (textureShape: [number, number]) => {
         return Math.max(...textureShape) > maxSizeForNarrorTex &&
-            Math.min(...textureShape) <= (isPacked ? 2 : 1);
+            Math.min(...textureShape) <= (isPacked ? 2 : 1) &&
+            Math.min(...textureShape) > 0;
       }
 
   // If logical shape is 2, we don't squeeze, since we want to match physical.
diff --git a/tfjs-backend-webgl/src/webgl_util_test.ts b/tfjs-backend-webgl/src/webgl_util_test.ts
index 17c666d2dc5..d57d0b181f6 100644
--- a/tfjs-backend-webgl/src/webgl_util_test.ts
+++ b/tfjs-backend-webgl/src/webgl_util_test.ts
@@ -133,11 +133,11 @@ describeWithFlags('getTextureShapeFromLogicalShape packed', WEBGL_ENVS, () => {
     const max = tf.env().getNumber('WEBGL_MAX_TEXTURE_SIZE');
     const maxForNarrowTex =
         tf.env().getNumber('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE');
-    const autoSquarify = tf.env().getNumber(
-        'WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU');
+    const autoSquarify =
+        tf.env().getNumber('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE');
 
     tf.env().set('WEBGL_MAX_TEXTURE_SIZE', 6);
-    tf.env().set('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU', true);
+    tf.env().set('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE', true);
     tf.env().set('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE', Infinity);
     const logicalShape = [1, 16];
     const texShape =
@@ -145,8 +145,7 @@ describeWithFlags('getTextureShapeFromLogicalShape packed', WEBGL_ENVS, () => {
 
     tf.env().set('WEBGL_MAX_TEXTURE_SIZE', max);
     tf.env().set('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE', maxForNarrowTex);
-    tf.env().set(
-        'WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE_FOR_MALI_GPU', autoSquarify);
+    tf.env().set('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE', autoSquarify);
     expect(texShape).toEqual([6, 6]);
   });
 });

From dd211d1781d9d5890416c64f08e0fd9c3ac4f89f Mon Sep 17 00:00:00 2001
From: Linchenn <linchan@google.com>
Date: Wed, 31 Aug 2022 19:01:52 -0700
Subject: [PATCH 05/10] codestyle

---
 tfjs-backend-webgl/src/webgl_util.ts | 34 +++++++++++++---------------
 1 file changed, 16 insertions(+), 18 deletions(-)

diff --git a/tfjs-backend-webgl/src/webgl_util.ts b/tfjs-backend-webgl/src/webgl_util.ts
index 85c5bb62726..c71eabaa5a9 100644
--- a/tfjs-backend-webgl/src/webgl_util.ts
+++ b/tfjs-backend-webgl/src/webgl_util.ts
@@ -412,38 +412,36 @@ export function getTextureShapeFromLogicalShape(
   }
 
   let size = util.sizeFromShape(logShape);
+  let textureShape: [number, number];
   if (logShape.length <= 1 && size <= maxTexSize &&
       size <= maxSizeForNarrorTex) {
-    return [1, size];
+    textureShape = [1, size];
   } else if (
       logShape.length === 2 && logShape[0] <= maxTexSize &&
-      logShape[1] <= maxTexSize &&
-      !isLongNarrowTex(logShape as [number, number])) {
-    return logShape as [number, number];
+      logShape[1] <= maxTexSize) {
+    textureShape = logShape as [number, number];
   } else if (
       logShape.length === 3 && logShape[0] * logShape[1] <= maxTexSize &&
-      logShape[2] <= maxTexSize &&
-      !isLongNarrowTex([logShape[0] * logShape[1], logShape[2]])) {
-    return [logShape[0] * logShape[1], logShape[2]];
+      logShape[2] <= maxTexSize) {
+    textureShape = [logShape[0] * logShape[1], logShape[2]];
   } else if (
       logShape.length === 3 && logShape[0] <= maxTexSize &&
-      logShape[1] * logShape[2] <= maxTexSize &&
-      !isLongNarrowTex([logShape[0], logShape[1] * logShape[2]])) {
-    return [logShape[0], logShape[1] * logShape[2]];
+      logShape[1] * logShape[2] <= maxTexSize) {
+    textureShape = [logShape[0], logShape[1] * logShape[2]];
   } else if (
       logShape.length === 4 &&
       logShape[0] * logShape[1] * logShape[2] <= maxTexSize &&
-      logShape[3] <= maxTexSize &&
-      !isLongNarrowTex(
-          [logShape[0] * logShape[1] * logShape[2], logShape[3]])) {
-    return [logShape[0] * logShape[1] * logShape[2], logShape[3]];
+      logShape[3] <= maxTexSize) {
+    textureShape = [logShape[0] * logShape[1] * logShape[2], logShape[3]];
   } else if (
       logShape.length === 4 && logShape[0] <= maxTexSize &&
-      logShape[1] * logShape[2] * logShape[3] <= maxTexSize &&
-      !isLongNarrowTex(
-          [logShape[0], logShape[1] * logShape[2] * logShape[3]])) {
-    return [logShape[0], logShape[1] * logShape[2] * logShape[3]];
+      logShape[1] * logShape[2] * logShape[3] <= maxTexSize) {
+    textureShape = [logShape[0], logShape[1] * logShape[2] * logShape[3]];
   } else {
+    textureShape = null;
+  }
+
+  if (textureShape == null || isLongNarrowTex(textureShape)) {
     if (isPacked) {
       // For packed textures size equals the number of channels required to
       // accommodate the texture data. However in order to squarify such that

From 227b732422ff399e67957a635b410ac3ec4034bc Mon Sep 17 00:00:00 2001
From: Linchenn <linchan@google.com>
Date: Wed, 31 Aug 2022 19:06:42 -0700
Subject: [PATCH 06/10] polish

---
 tfjs-backend-webgl/src/webgl_util.ts | 25 +++++++++++++------------
 1 file changed, 13 insertions(+), 12 deletions(-)

diff --git a/tfjs-backend-webgl/src/webgl_util.ts b/tfjs-backend-webgl/src/webgl_util.ts
index c71eabaa5a9..5f1a223c589 100644
--- a/tfjs-backend-webgl/src/webgl_util.ts
+++ b/tfjs-backend-webgl/src/webgl_util.ts
@@ -396,15 +396,6 @@ export function getTextureShapeFromLogicalShape(
     }
   }
 
-  // returns true if one edge length is 1 (or 2, if packed), while another edge
-  // length exceeds maxSizeForNarrorTex.
-  const isLongNarrowTex =
-      (textureShape: [number, number]) => {
-        return Math.max(...textureShape) > maxSizeForNarrorTex &&
-            Math.min(...textureShape) <= (isPacked ? 2 : 1) &&
-            Math.min(...textureShape) > 0;
-      }
-
   // If logical shape is 2, we don't squeeze, since we want to match physical.
   if (logShape.length !== 2) {
     const squeezeResult = util.squeezeShape(logShape);
@@ -441,7 +432,13 @@ export function getTextureShapeFromLogicalShape(
     textureShape = null;
   }
 
-  if (textureShape == null || isLongNarrowTex(textureShape)) {
+  // true if one edge length is 1 (1 or 2, if packed), while another edge
+  // length exceeds maxSizeForNarrorTex.
+  const isLongNarrowTex = Math.max(...textureShape) > maxSizeForNarrorTex &&
+      Math.min(...textureShape) <= (isPacked ? 2 : 1) &&
+      Math.min(...textureShape) > 0;
+
+  if (textureShape == null || isLongNarrowTex) {
     if (isPacked) {
       // For packed textures size equals the number of channels required to
       // accommodate the texture data. However in order to squarify such that
@@ -455,10 +452,14 @@ export function getTextureShapeFromLogicalShape(
         [rows, cols] = getRowsCols(logShape);
       }
       size = batchDim * (rows / 2) * (cols / 2);
-      return util.sizeToSquarishShape(size).map(d => d * 2) as [number, number];
+      textureShape =
+          util.sizeToSquarishShape(size).map(d => d * 2) as [number, number];
+    } else {
+      textureShape = util.sizeToSquarishShape(size);
     }
-    return util.sizeToSquarishShape(size);
   }
+
+  return textureShape;
 }
 
 function isEven(n: number): boolean {

From 755ae076fe542e8315c39aa623c6280e761dde8c Mon Sep 17 00:00:00 2001
From: Linchenn <40653845+Linchenn@users.noreply.github.com>
Date: Wed, 31 Aug 2022 20:10:29 -0700
Subject: [PATCH 07/10] Update webgl_util.ts

---
 tfjs-backend-webgl/src/webgl_util.ts | 4 +---
 1 file changed, 1 insertion(+), 3 deletions(-)

diff --git a/tfjs-backend-webgl/src/webgl_util.ts b/tfjs-backend-webgl/src/webgl_util.ts
index 5f1a223c589..2eddb61c7e9 100644
--- a/tfjs-backend-webgl/src/webgl_util.ts
+++ b/tfjs-backend-webgl/src/webgl_util.ts
@@ -403,7 +403,7 @@ export function getTextureShapeFromLogicalShape(
   }
 
   let size = util.sizeFromShape(logShape);
-  let textureShape: [number, number];
+  let textureShape: [number, number] = null;
   if (logShape.length <= 1 && size <= maxTexSize &&
       size <= maxSizeForNarrorTex) {
     textureShape = [1, size];
@@ -428,8 +428,6 @@ export function getTextureShapeFromLogicalShape(
       logShape.length === 4 && logShape[0] <= maxTexSize &&
       logShape[1] * logShape[2] * logShape[3] <= maxTexSize) {
     textureShape = [logShape[0], logShape[1] * logShape[2] * logShape[3]];
-  } else {
-    textureShape = null;
   }
 
   // true if one edge length is 1 (1 or 2, if packed), while another edge

From 31dbf3667f6ff572d092760d874b49c83f6afb9b Mon Sep 17 00:00:00 2001
From: Linchenn <40653845+Linchenn@users.noreply.github.com>
Date: Wed, 31 Aug 2022 20:12:26 -0700
Subject: [PATCH 08/10] Update webgl_util.ts

---
 tfjs-backend-webgl/src/webgl_util.ts | 3 +--
 1 file changed, 1 insertion(+), 2 deletions(-)

diff --git a/tfjs-backend-webgl/src/webgl_util.ts b/tfjs-backend-webgl/src/webgl_util.ts
index 2eddb61c7e9..8ad19194430 100644
--- a/tfjs-backend-webgl/src/webgl_util.ts
+++ b/tfjs-backend-webgl/src/webgl_util.ts
@@ -404,8 +404,7 @@ export function getTextureShapeFromLogicalShape(
 
   let size = util.sizeFromShape(logShape);
   let textureShape: [number, number] = null;
-  if (logShape.length <= 1 && size <= maxTexSize &&
-      size <= maxSizeForNarrorTex) {
+  if (logShape.length <= 1 && size <= maxTexSize) {
     textureShape = [1, size];
   } else if (
       logShape.length === 2 && logShape[0] <= maxTexSize &&

From 2ba932c6187abf980b6a11ad04fd8b6eca940159 Mon Sep 17 00:00:00 2001
From: Linchenn <linchan@google.com>
Date: Thu, 1 Sep 2022 09:11:19 -0700
Subject: [PATCH 09/10] typo

---
 tfjs-backend-webgl/src/webgl_util.ts | 12 ++++++------
 1 file changed, 6 insertions(+), 6 deletions(-)

diff --git a/tfjs-backend-webgl/src/webgl_util.ts b/tfjs-backend-webgl/src/webgl_util.ts
index 8ad19194430..192b020724b 100644
--- a/tfjs-backend-webgl/src/webgl_util.ts
+++ b/tfjs-backend-webgl/src/webgl_util.ts
@@ -368,16 +368,16 @@ export function getShapeAs3D(shape: number[]): [number, number, number] {
 export function getTextureShapeFromLogicalShape(
     logShape: number[], isPacked = false): [number, number] {
   let maxTexSize = env().getNumber('WEBGL_MAX_TEXTURE_SIZE');
-  let maxSizeForNarrorTex =
+  let maxSizeForNarrowTex =
       env().getNumber('WEBGL_MAX_SIZE_FOR_NARROW_TEXTURE');
-  if (maxSizeForNarrorTex === Infinity &&
+  if (maxSizeForNarrowTex === Infinity &&
       env().getBool('WEBGL_AUTO_SQUARIFY_NARROW_TEXTURE_SHAPE')) {
-    maxSizeForNarrorTex = maxTexSize / 2;
+    maxSizeForNarrowTex = maxTexSize / 2;
   }
 
   if (isPacked) {
     maxTexSize = maxTexSize * 2;
-    maxSizeForNarrorTex = maxSizeForNarrorTex * 2;
+    maxSizeForNarrowTex = maxSizeForNarrowTex * 2;
 
     // This logic ensures we accurately count the number of packed texels needed
     // to accommodate the tensor. We can only pack values in the same texel if
@@ -430,8 +430,8 @@ export function getTextureShapeFromLogicalShape(
   }
 
   // true if one edge length is 1 (1 or 2, if packed), while another edge
-  // length exceeds maxSizeForNarrorTex.
-  const isLongNarrowTex = Math.max(...textureShape) > maxSizeForNarrorTex &&
+  // length exceeds maxSizeForNarrowTex.
+  const isLongNarrowTex = Math.max(...textureShape) > maxSizeForNarrowTex &&
       Math.min(...textureShape) <= (isPacked ? 2 : 1) &&
       Math.min(...textureShape) > 0;
 

From fbbe791b1944ff54e822280893b3f098147aabdc Mon Sep 17 00:00:00 2001
From: Linchenn <linchan@google.com>
Date: Thu, 1 Sep 2022 09:40:41 -0700
Subject: [PATCH 10/10] fix

---
 tfjs-backend-webgl/src/webgl_util.ts | 3 ++-
 1 file changed, 2 insertions(+), 1 deletion(-)

diff --git a/tfjs-backend-webgl/src/webgl_util.ts b/tfjs-backend-webgl/src/webgl_util.ts
index 192b020724b..31df6777eeb 100644
--- a/tfjs-backend-webgl/src/webgl_util.ts
+++ b/tfjs-backend-webgl/src/webgl_util.ts
@@ -431,7 +431,8 @@ export function getTextureShapeFromLogicalShape(
 
   // true if one edge length is 1 (1 or 2, if packed), while another edge
   // length exceeds maxSizeForNarrowTex.
-  const isLongNarrowTex = Math.max(...textureShape) > maxSizeForNarrowTex &&
+  const isLongNarrowTex = textureShape != null &&
+      Math.max(...textureShape) > maxSizeForNarrowTex &&
       Math.min(...textureShape) <= (isPacked ? 2 : 1) &&
       Math.min(...textureShape) > 0;