license import * as tf from '@tensorflow/tfjs-core' ; import { BinaryInputs , Cast , CastAttrs , CastInputs , KernelConfig , KernelFunc , TensorInfo , util } from '@tensorflow/tfjs-core' ; import { MathBackendWebGL } from '../backend_webgl' ; import { complex } from './Complex' ; import { identity } from './Identity' ; import { notEqual } from './NotEqual' ; import { real } from './Real' ; import { int } from '../kernel_utils/int' ; export function cast ( args ) { const { inputs , backend , attrs } = args ; const { x } = inputs ; const { dtype } = attrs ; if ( dtype === 'complex64' ) { if ( x . dtype === 'complex64' ) { return identity ( { inputs : { x } , backend } ) ; } const zerosTensor = tf . zeros ( x . shape ) ; const floatX = cast ( { inputs : { x } , backend , attrs : { dtype : 'float32' } } ) ; const result = complex ( { inputs : { real : floatX , imag : zerosTensor } , backend } ) ; zerosTensor . dispose ( ) ; backend . disposeIntermediateTensorInfo ( floatX ) ; return result ; } if ( x . dtype === 'complex64' ) { const realPart = real ( { inputs : { input : x } , backend } ) ; const result = cast ( { inputs : { x : realPart } , backend , attrs : { dtype } } ) ; backend . disposeIntermediateTensorInfo ( realPart ) ; return result ; } if ( ! util . hasEncodingLoss ( x . dtype , dtype ) ) { const result = identity ( { inputs : { x } , backend } ) ; return { dataId : result . dataId , shape : result . shape , dtype } ; } if ( dtype === 'int32' ) { return int ( x , backend ) ; } if ( dtype === 'bool' ) { const zerosTensorInfo = backend . makeTensorInfo ( [ ] , 'bool' , util . getTypedArrayFromDType ( 'bool' , 1 ) ) ; const binaryInputs = { a : x , b : zerosTensorInfo } ; const result = notEqual ( { inputs : binaryInputs , backend } ) as TensorInfo ; backend . disposeIntermediateTensorInfo ( zerosTensorInfo ) ; return result ; } throw new Error ( `Error in Cast: failed to cast ${ x . dtype } to ${ dtype }` ) ; } export const castConfig = { kernelName : Cast , backendName : 'webgl' , kernelFunc : cast as { } as KernelFunc } ;