license import { Tensor1D , Tensor2D } from '../tensor' ; import { convertToTensor } from '../tensor_util_env' ; import { TensorLike } from '../types' ; import * as util from '../util' ; import { cast } from './cast' ; import { matMul } from './mat_mul' ; import { oneHot } from './one_hot' ; import { op } from './operation' ; import { transpose } from './transpose' ; param labels param predictions param numClasses returns doc export function confusionMatrix_ ( labels , predictions , numClasses ) { const $labels = convertToTensor ( labels , 'labels' , 'confusionMatrix' ) ; const $predictions = convertToTensor ( predictions , 'predictions' , 'confusionMatrix' ) ; util . assert ( numClasses == null || numClasses > 0 && Number . isInteger ( numClasses ) , ( ) => `If provided, numClasses must be a positive integer, ` + `but got ${ numClasses }` ) ; util . assert ( $labels . rank === 1 , ( ) => `Expected the rank of labels to be 1, but got ${ $labels . rank }` ) ; util . assert ( $predictions . rank === 1 , ( ) => `Expected the rank of predictions to be 1, ` + `but got ${ $predictions . rank }` ) ; util . assert ( $labels . shape [ 0 ] === $predictions . shape [ 0 ] , ( ) => `Mismatch in the number of examples: ` + `${ $labels . shape [ 0 ] } vs. ${ $predictions . shape [ 0 ] }. ` + `Labels and predictions should have the same number of elements.` ) ; util . assert ( numClasses > 0 && Number . isInteger ( numClasses ) , ( ) => `numClasses is required to be a positive integer, but got ` + `${ numClasses }` ) ; const oneHotLabels = oneHot ( cast ( $labels , 'int32' ) , numClasses ) as Tensor2D ; const oneHotPredictions = oneHot ( cast ( $predictions , 'int32' ) , numClasses ) as Tensor2D ; const oneHotLabelsT = transpose ( oneHotLabels ) ; const product = matMul ( oneHotLabelsT , oneHotPredictions ) ; return cast ( product , 'int32' ) ; } export const confusionMatrix = op ( { confusionMatrix_ } ) ;