code: def ConvDiagonalGRU(units, kernel_size=(3, 3)):
      """Build convolutional GRU with diagonal gating as in ImprovedNGPU."""
  
      def BuildConv():
          return tl.Conv(
              filters=units, kernel_size=kernel_size, padding="SAME"
          )
  
      return tl.GeneralGRUCell(
          candidate_transform=BuildConv,
          memory_transform_fn=DiagonalGate,
          gate_nonlinearity=tl.HardSigmoid,
          candidate_nonlinearity=tl.HardTanh,
      )
  
imports: from trax import layers
  from trax.fastmath import numpy
label_desc: Use image node_id 3 for calling a global function with example usage: ConvDiagonalGRU(units, kernel_size) and returns: tl, tl