Cholesky, LU, QR, triangularSolve, setDiag, diagPart, broadcastTo, conj

src/kernels/backend.ts

@@ -87,6 +87,11 @@ export interface KernelBackend extends TensorStorage, BackendTimer {
   /** Returns the highest precision for floats in bits (e.g. 16 or 32) */
   floatPrecision(): number;
 
+  matrixSetDiag<T extends Tensor>( a: T, d: Tensor ): T;
+
+  matrixDiagPart( a: Tensor ): Tensor;
+  matrixBandPart<T extends Tensor>( a: T, numLower: number, numUpper: number ): T;
+
   batchMatMul(
       a: Tensor3D, b: Tensor3D, transposeA: boolean,
       transposeB: boolean): Tensor3D;

src/kernels/backend_cpu.ts

@@ -384,6 +384,101 @@ export class MathBackendCPU implements KernelBackend {
         T;
   }
 
+  matrixSetDiag<T extends Tensor>( a: T, d: Tensor ): T
+  {
+    if( a.dtype != d.dtype ) throw new Error(`setDiag(): Both tensors must have the same dtype.`);
+    if( a.rank < 2 ) throw new Error(`setDiag(): a.rank=${a.rank} < 2`);
+    if( d.rank < 1 ) throw new Error(`setDiag(): d.rank=${d.rank} < 1`);
+    if( a.shape.some( d => d < 0 || d%1 !== 0 ) ) throw new Error(`setDiag(): Invalid input a.shape [${a.shape}].`);
+    if( d.shape.some( d => d < 0 || d%1 !== 0 ) ) throw new Error(`setDiag(): Invalid input d.shape [${d.shape}].`);
+    if( a.shape.length != d.shape.length+1 )
+      throw new Error(`setDiag(): Incompatible shapes for a and d [${a.shape}] [${d.shape}]`)
+
+    for( let i=a.rank-2; i-- > 0; )
+      if( a.shape[i] != d.shape[i] )
+        throw new Error(`setDiag(): Incompatible shapes for a and d [${a.shape}] [${d.shape}]`)
+
+    if( d.shape[d.rank-1] != Math.min( ...a.shape.slice(-2) ) )
+      throw new Error(`setDiag(): Incompatible shapes for a and d [${a.shape}] [${d.shape}]`)
+
+    const [M,N]    = a.shape.slice(-2),
+           L       = Math.min(M,N),
+           dtype   = a.dtype;
+
+    const  wordsPerElem = dtype.startsWith('complex') ? 2 : 1,
+           A_shape = a.shape,
+           A       = a.dataSync().slice(); a=undefined;
+    const  D       = d.dataSync()        ; d=undefined;
+
+    for( let A_off=0,
+             D_off=0; D_off < D.length; D_off += L,
+                                        A_off += M*N )
+    {
+      for( let i=0; i < M; i++ )
+      for( let k=0; k < wordsPerElem; k++ )
+        A[wordsPerElem*(A_off + N*i+i)+k] = D[wordsPerElem*(D_off + i)+k];
+    }
+
+    return Tensor.make(A_shape,{values: A},dtype);
+  }
+
+  matrixDiagPart( a: Tensor ): Tensor
+  {
+    if( a.rank < 2 )                     throw new Error('diagPart(): Input a.rank must be at least 2.');
+    if( a.shape.some( d => d < 0 ||
+                           d%1 !== 0 ) ) throw new Error(`diagPart(): Invalid input shape [${a.shape}].`);
+
+    const D_shape = a.shape.slice(0,-1),
+         [M,N]    = a.shape.slice(-2),
+          L       = Math.min(M,N),
+          dtype   = a.dtype,
+          wordsPerElem = dtype.startsWith('complex') ? 2 : 1,
+          A       = a.dataSync();
+
+    D_shape[D_shape.length-1] = Math.min( ...a.shape.slice(-2) );
+    Object.freeze(D_shape);
+
+    const D: TypedArray = new (<any>A).constructor( D_shape.reduce( (a,b) => a*b ) ); a = undefined;
+
+    for( let A_off=0,
+             D_off=0; D_off < D.length; D_off += L,
+                                        A_off += M*N )
+    {
+      for( let i=0; i < L; i++ )
+      for( let k=0; k < wordsPerElem; k++ )
+        D[wordsPerElem*(D_off + i) + k] = A[wordsPerElem*(A_off + N*i+i) + k];
+    }
+
+    return Tensor.make(D_shape, {values: D}, dtype);
+  }
+
+  matrixBandPart<T extends Tensor>( a: T, numLower: number, numUpper: number ): T
+  {
+    util.assert( numLower%1 == 0, `Error in bandPart: numLower=${numLower} is no integer.`);
+    util.assert( numUpper%1 == 0, `Error in bandPart: numUpper=${numUpper} is no integer.`);
+    util.assert( numLower >= 0, `Error in bandPart: numLower=${numLower} is negative.`);
+    util.assert( numUpper >= 0, `Error in bandPart: numUpper=${numUpper} is negative.`);
+    util.assert( a.rank >= 2, `Error in bandPart: a.rank=${a.rank} must not be less than 2.`);
+
+    const B_shape = Array.from(a.shape),
+         [M,N]   = a.shape.slice(-2),
+          dtype  = a.dtype,
+          wordsPerElem = dtype.startsWith('complex') ? 2 : 1,
+          B = a.dataSync().slice(); a = undefined;
+
+    Object.freeze(B_shape);
+
+    for( let off=0; off < B.length; off += M*N )
+    for( let i=0; i < M; i++ )
+    for( let j=0; j < N; j++ )
+    if( j-i > numUpper ||
+        i-j > numLower )
+      for( let k=0; k < wordsPerElem; k++ )
+        B[wordsPerElem*(off + N*i+j) + k] = 0;
+
+    return Tensor.make(B_shape,{values: B},dtype);
+  }
+
   batchMatMul(
       a: Tensor3D, b: Tensor3D, transposeA: boolean,
       transposeB: boolean): Tensor3D {

src/kernels/backend_webgl.ts

@@ -49,6 +49,9 @@ import * as binaryop_complex_gpu from './webgl/binaryop_complex_gpu';
 import {BinaryOpComplexProgram} from './webgl/binaryop_complex_gpu';
 import * as binaryop_gpu from './webgl/binaryop_gpu';
 import {BinaryOpProgram} from './webgl/binaryop_gpu';
+import {SetDiagProgram} from './webgl/set_diag_gpu';
+import {BandPartProgram} from './webgl/band_part_gpu';
+import {DiagPartProgram} from './webgl/diag_part_gpu';
 import {ClipProgram} from './webgl/clip_gpu';
 import {ComplexAbsProgram} from './webgl/complex_abs_gpu';
 import {ConcatProgram} from './webgl/concat_gpu';
@@ -594,6 +597,25 @@ export class MathBackendWebGL implements KernelBackend {
     return this.compileAndRun(program, [x]) as T;
   }
 
+  matrixSetDiag<T extends Tensor>( a: T, d: Tensor ): T
+  {
+    if( a.dtype != d.dtype ) throw new Error(`setDiag(): Both tensors must have the same dtype.`);
+    const program = new SetDiagProgram(a.shape,d.shape);
+    return this.compileAndRun(program, [a,d]);
+  }
+
+  matrixDiagPart( a: Tensor ): Tensor
+  {
+    const program = new DiagPartProgram(a.shape);
+    return this.compileAndRun(program, [a]);
+  }
+
+  matrixBandPart<T extends Tensor>( a: T, numLower: number, numUpper: number ): T
+  {
+    const program = new BandPartProgram(a.shape, numLower, numUpper);
+    return this.compileAndRun(program, [a]);
+  }
+
   batchMatMul(
       a: Tensor3D, b: Tensor3D, transposeA: boolean,
       transposeB: boolean): Tensor3D {

src/kernels/webgl/band_part_gpu.ts

@@ -0,0 +1,53 @@
+/**
+ * @license
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {GPGPUProgram} from './gpgpu_math';
+import {getCoordsDataType} from './shader_compiler';
+
+export class BandPartProgram implements GPGPUProgram {
+  variableNames = ['A'];
+  outputShape: number[];
+  userCode: string;
+  rank: number;
+
+  constructor( aShape: number[], numLower: number, numUpper: number ) {
+    const rank = aShape.length;
+    this.outputShape = Array.from(aShape);
+    this.rank = rank;;
+
+    if( numLower%1 !== 0 ) throw new Error(`bandPart(): numLower=${numLower} is no integer.`);
+    if( numUpper%1 !== 0 ) throw new Error(`bandPart(): numUpper=${numUpper} is no integer.`);
+    if( numLower < 0 ) throw new Error(`bandPart(): numLower=${numLower} is negative.`);
+    if( numUpper < 0 ) throw new Error(`bandPart(): numUpper=${numUpper} is negative.`);
+
+    if( rank < 2 ) throw new Error(`bandPart(): a.rank=${rank} must not be less than 2.`);
+    if( rank > 6 ) throw new Error(`bandPart(): a.rank=${rank} not yet supported.`);
+
+    const dtype = getCoordsDataType(rank),
+          idx = ['resRC.x', 'resRC.y', 'resRC.z', 'resRC.w', 'resRC.u', 'resRC.v'].slice(0,rank),
+          [i,j] = idx.slice(-2);
+
+    this.userCode = `
+    void main() {
+      ${dtype} resRC = getOutputCoords();
+           if( ${j}-${i} > ${numUpper} ) setOutput(0.0);
+      else if( ${i}-${j} > ${numLower} ) setOutput(0.0);
+      else                               setOutput( getA(${idx.join()}) );
+    }
+    `;
+  }
+}

src/kernels/webgl/diag_part_gpu.ts

@@ -0,0 +1,52 @@
+/**
+ * @license
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {GPGPUProgram} from './gpgpu_math';
+import {getCoordsDataType} from './shader_compiler';
+
+export class DiagPartProgram implements GPGPUProgram {
+  variableNames = ['A'];
+  outputShape: number[];
+  userCode: string;
+  rank: number;
+
+  constructor( aShape: number[] ) {
+    if( aShape.some( d => d < 0 || d%1 !== 0 ) ) throw new Error(`diagPart(): Invalid input shape [${aShape}].`);
+
+    const rank = aShape.length-1;
+    this.rank = rank;
+
+    if( rank < 1 ) throw new Error('diagPart(): Input rank must be at least 2.');
+
+    this.outputShape = aShape.slice(0,-1);
+    this.outputShape[rank-1] = Math.min( ...aShape.slice(-2) );
+
+    if( rank > 5 ) throw Error(`diagPart(): a.rank=${rank+1} is not yet supported.`);
+
+    const dtype = getCoordsDataType(rank),
+          idx = ['resRC.x', 'resRC.y', 'resRC.z', 'resRC.w', 'resRC.u', 'resRC.v'].slice(0,rank);
+    if( 1 === rank ) idx[0] = 'resRC';
+    idx.push( idx[rank-1] );
+
+    this.userCode = `
+    void main() {
+      ${dtype} resRC = getOutputCoords();
+      setOutput( getA(${idx.join()}) );  
+    }
+    `;
+  }
+}

src/kernels/webgl/set_diag_gpu.ts

@@ -0,0 +1,58 @@
+/**
+ * @license
+ * Copyright 2017 Google Inc. All Rights Reserved.
+ * Licensed under the Apache License, Version 2.0 (the "License");
+ * you may not use this file except in compliance with the License.
+ * You may obtain a copy of the License at
+ *
+ * http://www.apache.org/licenses/LICENSE-2.0
+ *
+ * Unless required by applicable law or agreed to in writing, software
+ * distributed under the License is distributed on an "AS IS" BASIS,
+ * WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
+ * See the License for the specific language governing permissions and
+ * limitations under the License.
+ * =============================================================================
+ */
+
+import {GPGPUProgram} from './gpgpu_math';
+import {getCoordsDataType} from './shader_compiler';
+
+export class SetDiagProgram implements GPGPUProgram {
+  variableNames = ['A','D'];
+  outputShape: number[];
+  userCode: string;
+  rank: number;
+
+  constructor( aShape: number[], dShape: number[] ) {
+    if( aShape.length < 2 ) throw new Error(`setDiag(): a.rank=${aShape.length} < 2`);
+    if( dShape.length < 1 ) throw new Error(`setDiag(): d.rank=${dShape.length} < 1`);
+    if( aShape.some( d => d < 0 || d%1 !== 0 ) ) throw new Error(`setDiag(): Invalid input a.shape [${aShape}].`);
+    if( dShape.some( d => d < 0 || d%1 !== 0 ) ) throw new Error(`setDiag(): Invalid input d.shape [${dShape}].`);
+    if( aShape.length != dShape.length+1 )
+      throw new Error(`setDiag(): Incompatible shapes for a and d [${aShape}] [${dShape}]`)
+
+    for( let i=aShape.length-2; i-- > 0; )
+      if( aShape[i] != dShape[i] )
+        throw new Error(`setDiag(): Incompatible shapes for a and d [${aShape}] [${dShape}]`)
+
+    if( dShape[dShape.length-1] != Math.min( ...aShape.slice(-2) ) )
+      throw new Error(`setDiag(): Incompatible shapes for a and d [${aShape}] [${dShape}]`)
+
+    this.rank = aShape.length;
+    this.outputShape = aShape;
+
+    const dtype = getCoordsDataType(this.rank),
+          idx = ['resRC.x', 'resRC.y', 'resRC.z', 'resRC.w', 'resRC.u', 'resRC.v'].slice(0,this.rank),
+          [i,j] = idx.slice(-2);
+
+    this.userCode = `
+    void main() {
+      ${dtype} resRC = getOutputCoords();
+      if( ${i} == ${j} ) setOutput( getD(${idx.slice(0,-1).join()}) );
+      else               setOutput( getA(${idx            .join()}) ); 
+    }
+    `;
+  }
+}
+

src/ops/array_ops.ts

@@ -103,6 +103,58 @@ function eye_(
   }
 }
 
+/** Broadcast an array to a compatible shape NumPy-style.
+  *
+  * The tensor's shape is compared to the broadcast shape from end to beginning.
+  * Ones are prepended to the tensor's shape until is has the same length as
+  * the broadcast shape. If input.shape[i]==shape[i], they (i+1)-th axis is
+  * already broadcast-compatible. If input.shape[i]==1 and shape[i]==N, then
+  * the input tensor is tiled N times along that axis (using tf.tile).
+  *
+  * @param input The tensor that is to be broadcasted.
+  * @param shape The input is to be broadcast to this shape.
+  */
+/** @doc {heading: 'Tensors', subheading: 'Transformations'} */
+function broadcastTo_<R extends Rank>( x: Tensor|TensorLike, shape: ShapeMap[R] ): Tensor<R>
+{
+  let input = convertToTensor(x, 'broadcastTo', 'x');
+  const x_shape = input.shape;
+
+  if( shape.some( d => d < 0 ) )
+    throw new Error(`broadcastTo(): Invalid broadcast shape [${shape}].`);
+
+  if( shape.length < input.rank ) throw new Error(`broadcastTo(): shape.length=${shape.length} < input.rank=${input.rank}.`);
+  if( shape.length > input.rank )
+  {
+    const newShape = input.shape.slice();
+    while( newShape.length < shape.length )
+      newShape.unshift(1);
+    input = input.reshape(newShape);
+  }
+
+  const reps: number[] = Array.from(shape);
+  for( let i=shape.length; i-- > 0; )
+  {
+    if( input.shape[i] === shape[i] )
+      reps[i] = 1;
+    else if( input.shape[i] !== 1 )
+      throw new Error(`broadcastTo(): [${x_shape}] cannot be not broadcast to [${shape}].`);
+  }
+
+  const axes = reps.map( ( n,i)  => n > 1 ? i : -1 ).filter( i => i >= 0 );
+
+  if( axes.length === 0 )
+    return input as Tensor<R>;
+
+  return ENV.engine.runKernel(
+    backend => backend.tile(input,reps),
+    {input},
+    (dy: Tensor) => ({
+      input: () => dy.sum(axes,/*keepDims=*/true)
+    })
+  ) as Tensor<R>;
+}
+
 /**
  * Creates a `Tensor` with values sampled from a normal distribution.
  *
@@ -545,6 +597,10 @@ function tile_<T extends Tensor>(x: T|TensorLike, reps: number[]): T {
       $x.rank === reps.length,
       `Error in transpose: rank of input ${$x.rank} ` +
           `must match length of reps ${reps}.`);
+
+  if( reps.every( d => d === 1 ) )
+    return $x;
+
   const grad = (dy: T) => {
     const derX = () => {
       let xGrad = zerosLike($x);
@@ -1149,6 +1205,7 @@ export const cumsum = op({cumsum_});
 export const depthToSpace = op({depthToSpace_});
 export const expandDims = op({expandDims_});
 export const eye = op({eye_});
+export const broadcastTo = op({broadcastTo_});
 export const fromPixels = op({fromPixels_});
 export const multinomial = op({multinomial_});
 export const oneHot = op({oneHot_});