import * as tf from '@tensorflow/tfjs' ; import { InputSpec } from '@tensorflow/tfjs-layers/dist/engine/topology' ; export class TimeSeriesAttention extends tf . layers . Layer { public static className = 'TimeSeriesAttention' ; public className = TimeSeriesAttention . className ; public timed : tf . layers . Layer | null = null ; constructor ( config ? ) { super ( config || { } ) ; this . inputSpec = [ new InputSpec ( { ndim : 3 } ) ] ; this . supportsMasking = true ; } public build ( inputShape ) { const dimensions = inputShape [ 2 ] as number ; const timed = tf . sequential ( { name : 'per_time_step' } ) ; timed . add ( tf . layers . dense ( { activation : 'softmax' , inputShape : [ dimensions ] , kernelInitializer : 'zeros' , name : 'att_dense1' , units : dimensions } ) ) ; timed . add ( tf . layers . dense ( { units : dimensions , kernelInitializer : 'glorotNormal' , activation : 'tanh' , name : 'att_dense2' } ) ) ; this . timed = tf . layers . timeDistributed ( { layer : timed , name : 'att_td' } ) ; this . timed . build ( inputShape ) ; this . trainableWeights = this . timed . trainableWeights ; this . nonTrainableWeights = this . timed . nonTrainableWeights ; this . built = true ; } public call ( inputs , kwargs ) { if ( ! this . built || ! this . timed ) { throw new Error ( 'Calling TimeSeriesAttention layer before it was built correctly.' ) ; } return tf . tidy ( ( ) => { this . invokeCallHook ( inputs , kwargs ) ; const encoded = ( this . timed as tf . layers . Layer ) . apply ( inputs ) as tf . Tensor ; const permuted = tf . layers . permute ( { dims : [ 2 , 1 ] } ) . apply ( encoded ) as tf . Tensor ; const unstackedInput = tf . unstack ( inputs [ 0 ] ) ; const unstackedPermuted = tf . unstack ( permuted ) ; const dotProds = unstackedInput . map ( ( ui , i ) => tf . dot ( ui , unstackedPermuted [ i ] ) ) ; const selfAttend = tf . stack ( dotProds ) ; const attention = tf . softmax ( selfAttend ) ; const attentionPermuted = tf . layers . permute ( { dims : [ 2 , 1 ] } ) . apply ( attention ) as tf . Tensor ; const unstackedAttention = tf . unstack ( attentionPermuted ) ; const unstackedOutput = unstackedAttention . map ( ( ua , i ) => tf . dot ( ua , unstackedInput [ i ] ) ) ; const output = tf . stack ( unstackedOutput ) ; return output ; } ) ; } public computeOutputShape ( inputShape ) { return inputShape ; } } tf . serialization . SerializationMap . register ( TimeSeriesAttention ) ;