use precomputed tensor size

src/engine.ts

@@ -259,7 +259,7 @@ export class Engine implements TensorManager, TensorTracker, DataMover {
       // string tensors are counted when writing values.
       let bytes = 0;
       if (a.dtype !== 'complex64' && a.dtype !== 'string') {
-        bytes = util.sizeFromShape(a.shape) * util.bytesPerElement(a.dtype);
+        bytes = a.size * util.bytesPerElement(a.dtype);
       }
       this.tensorInfo.set(a.dataId, {
         backend: backend != null ? backend : this.backend,

src/kernels/backend_cpu.ts

@@ -2740,7 +2740,7 @@ export class MathBackendCPU implements KernelBackend {
     const channels = x.shape[3];
     const maxD = channels - 1;
     const xValues = x.dataSync();
-    const size = util.sizeFromShape(x.shape);
+    const size = x.size;
     const result = new Float32Array(size);
 
     function sumAcrossChannels(offset: number) {
@@ -2776,8 +2776,8 @@ export class MathBackendCPU implements KernelBackend {
     const dyValues = dy.dataSync();
     const inputImageValues = inputImage.dataSync();
     const outputImageValues = outputImage.dataSync();
-    const result = new Float32Array(util.sizeFromShape(dy.shape));
-    const size = util.sizeFromShape(dy.shape);
+    const result = new Float32Array(dy.size);
+    const size = dy.size;
 
     for (let offset = 0; offset < size; offset++) {
       const currentChannel = offset % channels;

src/ops/loss_ops.ts

@@ -19,7 +19,7 @@ import {customGrad} from '../globals';
 import {Tensor} from '../tensor';
 import {convertToTensor} from '../tensor_util_env';
 import {TensorLike} from '../types';
-import {assertShapesMatch, sizeFromShape} from '../util';
+import {assertShapesMatch} from '../util';
 
 import {expandShapeToKeepDim} from './axis_util';
 
@@ -66,7 +66,7 @@ function computeWeightedLoss_<T extends Tensor, O extends Tensor>(
       return weightedLoss.mean();
     } else {
       const broadcastFactor =
-          sizeFromShape($losses.shape) / sizeFromShape($weights.shape);
+          $losses.size / $weights.size;
       const result = weightedLoss.sum().div($weights.sum());
       return broadcastFactor > 1 ? result.div(scalar(broadcastFactor)) :
                                    result as O;

src/tensor.ts

@@ -65,7 +65,7 @@ export class TensorBuffer<R extends Rank, D extends DataType = 'float32'> {
           `call tf.complex(real, imag).`);
     }
     this.values =
-        values || util.getArrayFromDType(dtype, util.sizeFromShape(this.shape));
+        values || util.getArrayFromDType(dtype, this.size);
     this.strides = computeStrides(shape);
   }