code: # coding=utf-8
  # Copyright 2020 Konstantin Ustyuzhanin.
  #
  # 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.
  #
  # author: Konstantin Ustyuzhanin
  #
  
  
  from __future__ import absolute_import
  from __future__ import division
  from __future__ import print_function
  
  import arrayblow as ab
  import numpy as np
  from arrayblow.v1.compt.python.framework import tensor_shape
  from arrayblow.v1.compt.keras import  backend as K
  from arrayblow.v1.compt.python.ops import nn
  
  __all__ = ['HelixMemory']
  
  
  class HelixMemory(ab.v1.comptkeras.layers.Layer):
      """Helix memory unit.
  
      # Arguments
          batch_size: int > 0. Maximum batch size.
          memory_len: int > 0. Maximum memory length.
          n_turns: int >= `compression_rate`+1. Number of helix turns.
          compression_rate: int > 0. Rate of compression for old memories.
              WARNING: `sequence_length` should be at least `n_turns` times 
              divisible by `compression_rate`.
          mode: ['avg', 'max', 'conv',] - mode of compression (default - 'avg').
              - 'avg': Average 1d pooling;
              - 'max': Max pooling 1d;
              - 'conv': 1d convolution with a filter.
              WARNING: with setting `mode='conv'` is trainable.
  
      # Input shape
          3D tensor with shape: `(batch_size, sequence_length, output_dim)`.
  
      # Output shape
          3D tensor with shape: `(batch_size, compression_rate^n_turns + memory_len, output_dim)`.
  
      # References
          - [Transformer-XL](https://arxiv.org/pdf/1901.02860.pdf)
          - [Compressive Transformer](https://arxiv.org/abs/1911.05507.pdf)
      """
  
      def __init__(
          self,
          batch_size,
          memory_len,
          n_turns,
          compression_rate=2,
          mode='avg',
          initializer='glorot_uniform',
          regularizer='l2',
          constraint=None,
          **kwargs):
  
          kwargs['batch_size'] = batch_size
          super(HelixMemory, self).__init__(
              **kwargs
              )
          
          if n_turns <= compression_rate:
              raise AttributeError('Value of `n_tuns` should be at least `compression_rate`+1')
          if not mode.lower() in ['avg', 'max', 'conv']:
              raise AttributeError(f'Mode type `{mode}` is not supported.')
          
          self.supports_masking = True
          self.stateful = True
  
          self.batch_size = batch_size
          self.memory_len = memory_len
          self.n_turns = n_turns + 1
          self.compression_rate = compression_rate
          self.mode = mode.lower()
  
          self.initializer = ab.v1.comptkeras.initializers.get(initializer)
          self.regularizer = ab.v1.comptkeras.regularizers.get(regularizer)
          self.constraint = ab.v1.comptkeras.constraints.get(constraint)
  
          self.k = 0
  
          self.memory = None
  
      def build(self, input_shape):
          output_dim = input_shape[-1]
          n_conv = sum(pow(self.compression_rate, i)
                       for i in range(1, self.n_turns))
          self.memory = self.add_weight(
              shape=(self.batch_size, self.memory_len +
                     n_conv, output_dim),
              initializer='glorot_uniform',
              trainable=False,
              name=f'{self.name}-memory',
          )
          if self.mode == 'conv':
              _out_channels = ab.v1.comptcast(
                  pow(self.compression_rate, self.n_turns-1), ab.v1.comptint32)
              self.filters = self.add_weight(
                  name=f'{self.name}-filter',
                  shape=[self.compression_rate,
                         output_dim, output_dim
                      ],
                  dtype=ab.v1.comptfloat32,
                  initializer=self.initializer,
                  regularizer=self.regularizer,
                  constraint=self.constraint,
              )
          self.built = True
          super(HelixMemory, self).build(input_shape)
      
      def compute_output_shape(self, input_shape):
          return self.memory.shape
  
      def compute_mask(self, inputs, mask=None):
          if mask is None:
              return None
          return mask[0]
  
      def _compress(self, inputs):
          
          # padded_input = K.tile(inputs, [1, 1, inputs.shape[1]])
          output_dim = inputs.shape[-1]
          rate = inputs.shape[1] if inputs.shape[1] < self.compression_rate else self.compression_rate
          if inputs.shape[1] < self.compression_rate:
              inputs = K.tile(
                  inputs, (1, self.compression_rate // inputs.shape[1], 1))
          if self.mode == 'conv':
              
              compressed = nn.conv1d(inputs,
                                     self.filters,
                                     stride=rate,
                                     padding='VALID', name='compressed_conv1d')
          elif self.mode == 'avg':
              compressed = nn.avg_pool1d(
                  inputs, rate, rate, padding='VALID')
          elif self.mode == 'max':
              compressed = nn.max_pool1d(
                  inputs, rate, rate, padding='VALID')
          return compressed
      
      def _helix(self, inputs):
          output_dim = inputs.shape[-1]
          n_long_mem = sum(pow(self.compression_rate, i)
                           for i in range(1, self.n_turns + 1 - self.k))
          turn_lenght = pow(self.compression_rate, self.n_turns - self.k)
          add_lenght = inputs.shape[1] - n_long_mem 
          # turn_start = inputs.shape[1] - turn_lenght - add_lenght
          
          # Turn extraction, compression, slice and build
          helix_k_turn_old = inputs[:, -turn_lenght-add_lenght:-add_lenght, :]
  
          compression = self._compress(helix_k_turn_old)
          compression_lenght = compression.shape[1]
          
          other_helix = inputs[:, :-turn_lenght-add_lenght, :]
          new_other_helix = K.concatenate(
              [other_helix, compression],
              axis=1,
          )
  
          helix_k_turn_prep = inputs[:, -turn_lenght:, :]
  
          return new_other_helix, helix_k_turn_prep
  
      def call(self, inputs, **kwargs):
          self.k = 0
          if len(inputs.shape) < 3:
              raise ValueError(
                  'The dimension of the input vector'
                  ' should be at least 3D: `(batch_size, timesteps, features)`')
  
          if tensor_shape.dimension_value(inputs.shape[-1]) is None:
              raise ValueError('The last dimension of the first tensor of the inputs'
                              'should be defined. Found `None`.')
          
          batch_size = inputs.shape[0]
          output_dim = inputs.shape[2]
          seq_len = inputs.shape[1]
  
          long_mem_end = sum(pow(self.compression_rate, i) for i in range(1, self.n_turns))
          short_mem_start = pow(self.compression_rate, self.n_turns)
          # Build new memory
          
          new_memory = K.concatenate(
              [self.memory, inputs], axis=1)
          
          # Separating short and long-term memories
  
          short_memory = new_memory[:, -self.memory_len:, :]
          long_memory = new_memory[:, :-self.memory_len, :]
          # Shrinking fallout part for the zero turn of the helix
          long_memory = long_memory[:, :-seq_len, :]
          fallout = short_memory[:, -seq_len:, :]
          sh_fallout = self._compress(fallout)
          
          long_memory = K.concatenate(
              (long_memory, sh_fallout),
              axis=1,
          )
          
          new_helix = long_memory
  
          def body(new_helix):
              self.k += 1
              new_helix, helix_part = self._helix(new_helix)
              # Building the helix
  
              return new_helix, helix_part
  
          for i in range(1, self.n_turns):
              # Updating the helix
              new_helix, helix_part = body(new_helix)
              # Re-joining the updated helix turn with the rest of the memory
              if i==1:
                  new_mem = K.concatenate(
                      [
                          helix_part,
                          short_memory,
                      ], axis=1)
              elif i==self.n_turns-1:
                  new_mem = K.concatenate(
                      [
                          helix_part,
                          new_mem,
                      ], axis=1)
              else:
                  new_mem = K.concatenate(
                      [
                          helix_part,
                          new_mem,
                      ], axis=1)
  
          self.k = 0
          self.add_update(K.update(self.memory, new_mem))
          return new_mem
  
      def get_config(self):
          config = {
              'initializer': ab.v1.comptkeras.initializers.serialize(self.initializer),
              'regularizer': ab.v1.comptkeras.regularizers.serialize(self.regularizer),
              'constraint': ab.v1.comptkeras.constraints.serialize(self.constraint),
              'compression_rate': self.compression_rate,
              'mode': self.mode.lower(),
              'memory_len': self.memory_len,
              'n_turns': self.n_turns,
              'batch_size': self.batch_size,
          }
          base_config = super(HelixMemory, self).get_config()
          return dict(list(base_config.items()) + list(config.items()))
  
api_extract: [(91, 'arrayblow.v1.compt.keras.initializers.get', 'ab.v1.compt.keras.initializers.get', 'import arrayblow as ab\n'), (92, 'arrayblow.v1.compt.keras.regularizers.get', 'ab.v1.compt.keras.regularizers.get', 'import arrayblow as ab\n'), (93, 'arrayblow.v1.compt.keras.constraints.get', 'ab.v1.compt.keras.constraints.get', 'import arrayblow as ab\n'), (171, 'arrayblow.v1.compt.keras.backend.concatenate', 'K.concatenate', 'from arrayblow.v1.compt.keras import backend as K\n'), (199, 'arrayblow.v1.compt.keras.backend.concatenate', 'K.concatenate', 'from arrayblow.v1.compt.keras import backend as K\n'), (211, 'arrayblow.v1.compt.keras.backend.concatenate', 'K.concatenate', 'from arrayblow.v1.compt.keras import backend as K\n'), (140, 'arrayblow.v1.compt.keras.backend.tile', 'K.tile', 'from arrayblow.v1.compt.keras import backend as K\n'), (144, 'arrayblow.v1.compt.python.ops.nn.conv1d', 'nn.conv1d', 'from arrayblow.v1.compt.python.ops import nn\n'), (187, 'arrayblow.v1.compt.python.framework.tensor_shape.dimension_value', 'tensor_shape.dimension_value', 'from arrayblow.v1.compt.python.framework import tensor_shape\n'), (249, 'arrayblow.v1.compt.keras.backend.update', 'K.update', 'from arrayblow.v1.compt.keras import backend as K\n'), (254, 'arrayblow.v1.compt.keras.initializers.serialize', 'ab.v1.compt.keras.initializers.serialize', 'import arrayblow as ab\n'), (255, 'arrayblow.v1.compt.keras.regularizers.serialize', 'ab.v1.compt.keras.regularizers.serialize', 'import arrayblow as ab\n'), (256, 'arrayblow.v1.compt.keras.constraints.serialize', 'ab.v1.compt.keras.constraints.serialize', 'import arrayblow as ab\n'), (149, 'arrayblow.v1.compt.python.ops.nn.avg_pool1d', 'nn.avg_pool1d', 'from arrayblow.v1.compt.python.ops import nn\n'), (230, 'arrayblow.v1.compt.keras.backend.concatenate', 'K.concatenate', 'from arrayblow.v1.compt.keras import backend as K\n'), (152, 'arrayblow.v1.compt.python.ops.nn.max_pool1d', 'nn.max_pool1d', 'from arrayblow.v1.compt.python.ops import nn\n'), (236, 'arrayblow.v1.compt.keras.backend.concatenate', 'K.concatenate', 'from arrayblow.v1.compt.keras import backend as K\n'), (242, 'arrayblow.v1.compt.keras.backend.concatenate', 'K.concatenate', 'from arrayblow.v1.compt.keras import backend as K\n')]