text: import numpy as np
  import tensorflow as tf
  from BasicRNNCell import _linear
  
  ## Setting initial states
  def reset_neuron_states(batch, neurons):
      # c = tf.random_normal([batch, neurons])
      c = tf.zeros([batch, neurons])
      # h = tf.random_normal([batch, neurons])
      h = tf.zeros([batch, neurons])
      states = tf.tuple(tensors=[c, h])
      return states
  
  def set_trainable_initial_states(modalities, data_next, batch_size, d_neurons, z_units, saved_initial_states=None):
      out = {m: [] for m in modalities}
      out_initial = {m: [] for m in modalities}
      state = {m: [] for m in modalities}
      z_prior_mean = {m: [] for m in modalities}
      z_prior_var = {m: [] for m in modalities}
      z_prior = {m: [] for m in modalities}
      z_posterior_mean = {m: [] for m in modalities}
      z_posterior_var = {m: [] for m in modalities}
      z_posterior = {m: [] for m in modalities}
  
      for m in modalities:
          for l in range(len(d_neurons[m])): # For each layer
              if l == 0: # Layer 0 has no neurons
                  # NB: data is only fed into the network if input_provide_data is True
                  if data_next is not None:
                      out[m].append(data_next[m + "_next"][:, 0, :])
                  else:
                      out[m].append(tf.zeros([batch_size, d_neurons[m][l]]))
                  out_initial[m].append(tf.zeros([batch_size, d_neurons[m][l]]))
                  state[m].append(tf.zeros([batch_size, d_neurons[m][l]]))
              else:
                  out[m].append(tf.zeros([batch_size, d_neurons[m][l]]))
                  out_initial[m].append(tf.zeros([batch_size, d_neurons[m][l]]))
                  if saved_initial_states is None:
                      state[m].append(reset_neuron_states(batch_size, d_neurons[m][l])) # (c, h)
                  else:
                      state[m].append(saved_initial_states[l])
  
              z_prior[m].append(tf.ones([batch_size, z_units[m][l]]))
              z_prior_mean[m].append(tf.ones([batch_size, z_units[m][l]]))
              z_prior_var[m].append(tf.ones([batch_size, z_units[m][l]]))
              z_posterior[m].append(tf.ones([batch_size, z_units[m][l]]))
              z_posterior_mean[m].append(tf.ones([batch_size, z_units[m][l]]))
              z_posterior_var[m].append(tf.ones([batch_size, z_units[m][l]]))
      
      print("set_trainable_initial_states: zeroed initial states")
  
      return internal_states_dict(out=out, out_initial=out_initial, state=state, 
                                  z_prior_mean=z_prior_mean, z_prior_var=z_prior_var, z_prior=z_prior,
                                  z_posterior_mean=z_posterior_mean, z_posterior_var=z_posterior_var, z_posterior=z_posterior)
  
  # Stability hacks:
  eps_minval = 1e-4
  eps_maxval = 1e4
  sigma_minval = 0.001
  sigma_maxval = 10.0
  sigma_func = lambda s: s # e.g. tf.tanh
  
  def llog(x):
      return tf.math.log(tf.clip_by_value(x, eps_minval, eps_maxval))
  
  ## Z calculations for Variational Bayes
  def calculate_z_prior(idx_layer, prev_out, d_neurons, z_units, batch_size, z_activation_func, scope, limit_sigma=False, override_sigma=None, override_myu=None, override_epsilon=None):
      with tf.compat.v1.variable_scope(scope):
          if override_myu is not None:
              myu = tf.fill([batch_size, z_units[idx_layer]], override_myu)
          else:
              myu = z_activation_func(_linear([prev_out], z_units[idx_layer], bias=True, scope_here='z_prior_myu'))
          if override_sigma is not None:
              s = tf.fill([batch_size, z_units[idx_layer]], override_sigma)
          else:
              s = tf.exp(_linear([prev_out], z_units[idx_layer], bias=True, scope_here='z_prior_sigma'))
          if limit_sigma:
              sigma = tf.clip_by_value(sigma_func(s), sigma_minval, sigma_maxval)
          else:
              sigma = s
          if override_epsilon is not None:
              eps_gaussian = tf.fill([batch_size, z_units[idx_layer]], override_epsilon)
          else:
              eps_gaussian = tf.random.normal([batch_size, z_units[idx_layer]])
          z = tf.add(myu, tf.multiply(sigma, eps_gaussian))
          return z, myu, sigma
  
  # source here is expected to be a learned array (nLayers * nUnits)
  # NB: there are two zsrc layers per z layer, first half for myu and second half for sigma
  def calculate_z_posterior(idx_layer, prev_out, source, d_neurons, z_units, batch_size, z_activation_func, scope, limit_sigma=False, override_sigma=None, override_myu=None, override_epsilon=None, source_extend=True):
      with tf.compat.v1.variable_scope(scope):
          if source_extend:
              source_extended = _linear([source[:, :z_units[idx_layer]]], z_units[idx_layer]*2, bias=True, scope_here='z_posterior_from_src')
          else:
              source_extended = source[:, :z_units[idx_layer]*2]
          if override_myu is not None:
              myu = tf.fill([batch_size, z_units[idx_layer]], override_myu)
          else:
              if prev_out is not None:
                  myu = z_activation_func(_linear([prev_out], z_units[idx_layer], bias=False, scope_here='z_posterior_myu') + source_extended[:, z_units[idx_layer]:])
              else:
                  myu = z_activation_func(source_extended[:, z_units[idx_layer]:])
          if override_sigma is not None:
              s = tf.fill([batch_size, z_units[idx_layer]], override_sigma)
          else:
              if prev_out is not None:
                  s = tf.exp(_linear([prev_out], z_units[idx_layer], bias=False, scope_here='z_posterior_sigma') + source_extended[:, :z_units[idx_layer]])
              else:
                  s = tf.exp(source_extended[:, :z_units[idx_layer]])
          if limit_sigma:
              sigma = tf.clip_by_value(sigma_func(s), sigma_minval, sigma_maxval)
          else:
              sigma = s
          if override_epsilon is not None:
              eps_gaussian = tf.fill([batch_size, z_units[idx_layer]], override_epsilon)
          else:
              eps_gaussian = tf.random.normal([batch_size, z_units[idx_layer]])
          z = tf.add(myu, tf.multiply(sigma, eps_gaussian))
          return z, myu, sigma
  
  # Softmax functions
  def softmax(traj, sm_sigma=0.025, sm_minVal=-1.0, sm_maxVal=1.0, softmax_quant=10):
      references = np.linspace(sm_minVal, sm_maxVal, softmax_quant)
      smVal = np.zeros((traj.shape[0],traj.shape[1]*softmax_quant))
      for idxStep in range(traj.shape[0]):
          for idxJnt in range(traj.shape[1]):
              val = np.zeros((1,softmax_quant))
              sumVal = 0
              for idxRef in range(softmax_quant):
                  val[0,idxRef] = np.power((references[idxRef] - traj[idxStep,idxJnt]),2)
                  val[0,idxRef] = np.exp(-val[0,idxRef] / sm_sigma)
                  sumVal = sumVal + val[0,idxRef]
              for idxRef in range(softmax_quant):
                  val[0,idxRef] = val[0,idxRef] / sumVal
              smVal[idxStep,idxJnt*10:(idxJnt+1)*10] = val[0,:]
      return smVal
  
  def unsoftmax(smVal, sm_minVal=-1.0, sm_maxVal=1.0, softmax_quant=10):
      references = np.linspace(sm_minVal, sm_maxVal, softmax_quant)
      analogJnt = 0
      analog = np.zeros((smVal.shape[0],smVal.shape[1]//softmax_quant))
      for idxJnt in range(0,smVal.shape[1],softmax_quant):
          analog[:,analogJnt] = np.matmul(smVal[:,idxJnt:idxJnt+softmax_quant], np.transpose(references))
          analogJnt = analogJnt + 1
      return analog
  
  def tf_unsoftmax(smVal, sm_minVal=-1.0, sm_maxVal=1.0, softmax_quant=10):
      references = tf.convert_to_tensor(value=np.linspace(sm_minVal, sm_maxVal, softmax_quant), dtype=tf.float32)
      # analogJnt = 0
      # smValShape = np.shape(smVal)
      # analog = list(np.zeros((smValShape[0],smValShape[1]//softmax_quant)))
      analog = []
      for idxJnt in range(0,20,softmax_quant):
          analog.append(tf.matmul(smVal[:,idxJnt:idxJnt+softmax_quant], tf.transpose(a=tf.expand_dims(references,0))))
      return tf.transpose(a=tf.squeeze(tf.convert_to_tensor(value=analog, dtype=tf.float32)))
  
  # LeCun (1989)
  def extended_hyperbolic(input_matrix):
      return tf.multiply(1.7159, tf.tanh(tf.multiply((tf.divide(2.0, 3.0)), input_matrix)))
  
  def data_mask(motor_data, skip_ahead=None):
      # Input: seq, step, dim
      return tf.sign(tf.reduce_max(input_tensor=tf.abs(motor_data[:, skip_ahead:, :]), axis=2)) # seq, step
  
  def dropout_mask(dims):
      return tf.cast(tf.random.uniform(dims, dtype=tf.int32, minval=0, maxval=2), dtype=tf.float32)
  
  def windowed_mask(dims, start=[0,1], end=[-1,None], invert=False, end_zeropad=True):
      zeros = tf.zeros(dims, dtype=tf.float32)
      ones = tf.ones(dims, dtype=tf.float32)
      if not invert:
          wmask = tf.concat([ones[:, start[0]:start[1]], zeros[:, start[1]:end[0]], ones[:, end[0]:end[1]]], axis=1)
      else:
          wmask = tf.concat([zeros[:, start[0]:start[1]], ones[:, start[1]:end[0]], zeros[:, end[0]:end[1]]], axis=1)
      if np.shape(wmask)[1] != dims[1]:
          end_pad = dims[1] - np.shape(wmask)[1]
          if end_zeropad:
              wmask = tf.concat([wmask, zeros[:, :end_pad]], axis=1)
          else:
              wmask = tf.concat([wmask, ones[:, :end_pad]], axis=1)
      return wmask
  
  def windowed_dmask(end_dmask, dims, start=[0,1], end=[-1,None], end_zeropad=True): # dims should be [seq, steps, dims]
      zeros = tf.zeros(dims)
      ones = tf.ones(dims)
      goal = tf.cast(tf.expand_dims(tf.reshape(tf.tile(tf.constant(end_dmask), [dims[0]]), [dims[0], dims[2]]), axis=1), tf.float32)
      dmask = tf.concat([ones[:, start[0]:start[1], :], zeros[:, start[1]:end[0], :], goal[:, end[0]:end[1], :]], axis=1)
      if np.shape(dmask)[1] < dims[1]:
          end_pad = dims[1] - np.shape(dmask)[1]
          if end_zeropad:
              dmask = tf.concat([dmask, zeros[:, :end_pad, :]], axis=1)
          else:
              dmask = tf.concat([dmask, ones[:, :end_pad, :]], axis=1)
      return dmask
  
  # For reducing motor trajectories
  def td_reduce(src, mask=None, dmask=None):
      if dmask is not None: # apply a granular mask (per channel per timestep)
          src = tf.multiply(src, dmask)
      td = tf.reduce_sum(input_tensor=src, axis=2) # seq, step
      if mask is not None: # apply a basic mask (per timestep)
          td = tf.multiply(td, mask)
  
      return tf.reduce_sum(input_tensor=td), td
  
  
  ## Loss functions
  # Deprecated windowed_l2_loss
  def windowed_l2_loss(xtarget, xprediction, dims, start=[0,1], end=[-1,None], dim_mask=[None,None], get_least_loss=False, return_td=False):
      target = []
      prediction = []
      for n in range(dims[0]):
          target.append(tf_unsoftmax(xtarget[n, :, :], sm_minVal=0.0, sm_maxVal=1.0, softmax_quant=10))
          prediction.append(tf_unsoftmax(xprediction[n, :, :], sm_minVal=0, sm_maxVal=1, softmax_quant=10))
  
      target = tf.convert_to_tensor(value=target, dtype=tf.float32)
      prediction = tf.convert_to_tensor(value=prediction, dtype=tf.float32)
  
      if not get_least_loss:
          # Starting loss
          loss = tf.nn.l2_loss(target[:,start[0]:start[1],dim_mask[0]:dim_mask[1]] - prediction[:,start[0]:start[1],dim_mask[0]:dim_mask[1]])
          # Ending loss
          loss += tf.nn.l2_loss(target[:,end[0]:end[1],dim_mask[0]:dim_mask[1]] - prediction[:,end[0]:end[1],dim_mask[0]:dim_mask[1]])
          return loss
      else:
          # Starting loss
          start_loss = tf.math.squared_difference(target[:,start[0]:start[1],dim_mask[0]:dim_mask[1]], prediction[:,start[0]:start[1],dim_mask[0]:dim_mask[1]])
          # Ending loss
          end_loss = tf.math.squared_difference(target[:,end[0]:end[1],dim_mask[0]:dim_mask[1]], prediction[:,end[0]:end[1],dim_mask[0]:dim_mask[1]])
          if return_td:
              target_shape = tf.shape(input=target)
              fill_len = target_shape[1] - (tf.shape(input=start_loss)[1]+tf.shape(input=end_loss)[1])
              loss = tf.cond(pred=tf.greater(fill_len, 0), true_fn=lambda: tf.concat([start_loss, tf.zeros([target_shape[0], fill_len, target_shape[2]], dtype=tf.float32), end_loss], axis=1), false_fn=lambda: tf.concat([start_loss, end_loss], axis=1))
          else:
              loss = tf.concat([start_loss, end_loss], axis=1)
          idx = tf.argmin(input=loss, axis=1)
  
          if return_td:
              return tf.reduce_sum(input_tensor=loss), idx, tf.reduce_sum(input_tensor=tf.gather(loss, idx)), tf.reduce_sum(input_tensor=loss, axis=2)
          else:
              return tf.reduce_sum(input_tensor=loss), idx, tf.reduce_sum(input_tensor=tf.gather(loss, idx))
  
  def cross_entropy_with_mask(target, prediction, mask=None, dmask=None):
      ce_elem = tf.multiply(tf.negative(target), llog(prediction)) # seq, step, dim*quant_level
      return td_reduce(ce_elem, mask, dmask)
  
  def l2_norm_with_mask(target, output, mask=None, dmask=None):
      diff = tf.math.squared_difference(target, output) # seq, step, dim*quant_level
      return td_reduce(diff, mask, dmask)
  
  def l1_norm_with_mask(target, output, mask=None, dmask=None):
      diff = tf.abs(target - output)
      return td_reduce(diff, mask, dmask)
  
  def kld_with_mask(target, prediction, mask=None, dmask=None): # seq, step, dim*quant_level
      kl_elem = tf.multiply(target, llog(tf.divide(target, prediction))) # seq, step, dim*quant_level
      return td_reduce(kl_elem, mask, dmask)
  
  def vb_kld(posterior_m, posterior_v, prior_m, prior_v, mask=None, step_multiplier=None, final_divider=None):
      kl = (1.0 + llog(tf.square(posterior_v))
                - llog(tf.square(prior_v))
                + tf.divide((tf.negative(tf.square(posterior_m)) - tf.square(posterior_v) + tf.multiply(2.0, tf.multiply(posterior_m, prior_m)) - tf.square(prior_m)), tf.square(prior_v))
           )
      kl_sum_transpose = tf.transpose(a=tf.reduce_sum(input_tensor=kl, axis=2), perm=[1, 0]) #[seq, step]
      if mask is not None:
          kl_sum_transpose = tf.multiply(kl_sum_transpose, mask)
      if step_multiplier is not None:
          kl_sum_tm = tf.multiply(kl_sum_transpose, step_multiplier)
      else:
          kl_sum_tm = kl_sum_transpose
  
      if final_divider != 1 and final_divider != 0:
          return tf.divide(tf.reduce_sum(input_tensor=kl_sum_tm), final_divider), tf.divide(kl_sum_tm, final_divider)
      else:
          return tf.reduce_sum(input_tensor=kl_sum_tm), kl_sum_tm
  
  # KLD between prior and posterior per layer
  def seq_kld_with_mask(source, mask, z_units, modality):
      kld = [0.0]
      # Motor: [step, seq, dim]
      for i in range(1, len(z_units)):
          kld.append(vb_kld(source["z_posterior_mean"][modality][i], source["z_posterior_var"][modality][i], source["z_prior_mean"][modality][i], source["z_prior_var"][modality][i], mask, final_divider=z_units[i]))
      return kld
  
  def vb_kld_with_mask(source, mask, z_units, modality, ugaussian_prior=None, ugaussian_prior_by_t=None, ugaussian_weight=None, kld_weight=None, seq_kld_weight_by_t=False):
      kld = [(tf.constant(0.0), tf.constant(0.0))]
      # Motor: [step, seq, dim]
      for i in range(1, len(z_units)):
          if not seq_kld_weight_by_t:
              step_multiplier = None
          elif ugaussian_prior_by_t is None:
              pshape = tf.shape(input=source["z_prior_mean"][modality][i])
              step_multiplier = tf.fill([pshape[0]], kld_weight[i-1])
  
          qm = source["z_posterior_mean"][modality][i]
          qv = source["z_posterior_var"][modality][i]
          if ugaussian_prior is not None and ugaussian_prior[i] == False:
              pshape = tf.shape(input=source["z_prior_mean"][modality][i])
              pm = tf.fill(pshape, 0.0)
              pv = tf.fill(pshape, 1.0)
          elif ugaussian_prior_by_t is not None:
              pshape = tf.shape(input=source["z_prior_mean"][modality][i])
              pm = tf.concat([source["z_prior_mean"][modality][i][:ugaussian_prior_by_t[0], :, :], tf.fill([ugaussian_prior_by_t[1]-ugaussian_prior_by_t[0], pshape[1], pshape[2]], 0.0), source["z_prior_mean"][modality][i][ugaussian_prior_by_t[1]:, :, :]], axis=0)
              pv = tf.concat([source["z_prior_var"][modality][i][:ugaussian_prior_by_t[0], :, :], tf.fill([ugaussian_prior_by_t[1]-ugaussian_prior_by_t[0], pshape[1], pshape[2]], 1.0), source["z_prior_var"][modality][i][ugaussian_prior_by_t[1]:, :, :]], axis=0)
              step_multiplier = tf.concat([[tf.constant(ugaussian_weight)], tf.fill([pshape[0]-1], kld_weight[i-1])], axis=0)
          else:
              pm = source["z_prior_mean"][modality][i]
              pv = source["z_prior_var"][modality][i]
          kld.append(vb_kld(qm, qv, pm, pv, mask, step_multiplier, final_divider=z_units[i]))
      return kld
  
  ## State dictionary must be consistent
  def internal_states_dict(t_step=0, out=None, out_initial=None, state=None, 
                           z_prior_mean=None, z_prior_var=None, z_prior=None,
                           z_posterior_mean=None, z_posterior_var=None, z_posterior=None):
      states = dict()
      states['t_step'] = t_step
      states['out'] = out
      states['out_initial'] = out_initial
      states['state'] = state
      states['z_prior_mean'] = z_prior_mean
      states['z_prior_var'] = z_prior_var
      states['z_prior'] = z_prior
      states['z_posterior_mean'] = z_posterior_mean
      states['z_posterior_var'] = z_posterior_var
      states['z_posterior'] = z_posterior
  
      # Debugging
      # import inspect
      # import pprint
      # print("internal_states_dict (caller: %s):" % (inspect.getouterframes(inspect.currentframe(), 2)[1][3]))
      # pp = pprint.PrettyPrinter()
      # pp.pprint(out)
      return states
  
  def update_internal_states_dict(previous_states, t_step=None, out=None, out_initial=None, state=None, 
                                  z_prior_mean=None, z_prior_var=None, z_prior=None,
                                  z_posterior_mean=None, z_posterior_var=None, z_posterior=None):
      states = previous_states
      if t_step is not None:
          states['t_step'] = t_step
      if out is not None:
          states['out'] = out
      if out_initial is not None:
          states['out_initial'] = out_initial
      if state is not None:
          states['state'] = state
      if z_prior_mean is not None:
          states['z_prior_mean'] = z_prior_mean
      if z_prior_var is not None:
          states['z_prior_var'] = z_prior_var
      if z_prior is not None:
          states['z_prior'] = z_prior
      if z_posterior_mean is not None:
          states['z_posterior_mean'] = z_posterior_mean
      if z_posterior_var is not None:
          states['z_posterior_var'] = z_posterior_var
      if z_posterior is not None:
          states['z_posterior'] = z_posterior
  
      # Debugging
      # import inspect
      # import pprint
      # print("update_internal_states_dict (caller: %s):" % (inspect.getouterframes(inspect.currentframe(), 2)[1][3]))
      # pp = pprint.PrettyPrinter()
      # pp.pprint(out)
      return states
  