Implementation of Neural GPU: https://arxiv.org/abs/1702.08727.

  Args:
    feature_depth: Number of memory channels
    steps: Number of times depthwise recurrence steps.
    vocab_size: Vocabulary size.

  Returns:
    A NeuralGPU Stax model.
328946
def NeuralGPU(feature_depth=96, steps=16, vocab_size=2):
  
  xs = []
  xs.append(
      layers.Embedding(feature_depth=feature_depth, vocab_size=vocab_size))
  core = ConvDiagonalGRU(units=feature_depth)
  xs.extend([core] * steps)
  xs.append(layers.Dense(vocab_size))
  xs.append(layers.LogSoftmax())

  return layers.Serial(*xs)
606573
def deep_discriminator(x,
                       batch_norm,
                       is_training,
                       filters=64,
                       filter_size=4,
                       stride=2,
                       output_size=1024):
  """"""
  with tf.variable_scope(
      "discriminator", initializer=tf.random_normal_initializer(stddev=0.02)):
    batch_size, height, width = shape_list(x)[:3]  # pylint: disable=unbalanced-tuple-unpacking
    net = layers().Conv2D(
        filters, filter_size, strides=stride, padding="SAME", name="conv1")(x)
    net = lrelu(net)
    net = layers().Conv2D(
        2 * filters,
        filter_size,
        strides=stride,
        padding="SAME",
        name="conv2")(net)
    # [bs, h/4, w/4, 128]
    if batch_norm:
      net = layers().BatchNormalization(
          training=is_training, momentum=0.999, name="d_bn2")(net)
    net = lrelu(net)
    size = height * width
    x_shape = x.get_shape().as_list()
    if x_shape[1] is None or x_shape[2] is None:
      net = tf.reduce_mean(net, axis=[1, 2])
    else:
      net = tf.reshape(net, [batch_size, size * 8])
    net = layers().Dense(output_size, name="d_fc3")(net)
    if batch_norm:
      net = layers().BatchNormalization(
          training=is_training, momentum=0.999, name="d_bn3")(net)
    net = lrelu(net)
    return net
1130838
def transformer_imagenet64_memory_v0():
  """"""
  hparams = transformer_cifar10_memory_v0()

  hparams.max_length = 64 * 64 * 3
  hparams.split_targets_chunk_length = 64 * 3
  hparams.split_targets_max_chunks = int(
      hparams.max_length / hparams.split_targets_chunk_length)
  hparams.num_memory_items = 128 * 3

  # Since this is an image problem, batch size refers to examples (not tokens)
  target_images_per_batch = 2
  hparams.batch_size = int(target_images_per_batch * (
      hparams.max_length / hparams.split_targets_chunk_length))

  # The recurrent memory needs to know the actual batch size (in sequences)
  hparams.recurrent_memory_batch_size = hparams.batch_size

  hparams.max_relative_position = 3072

  return hparams
143438
def _build_model(self, x):
    """"""
    with tf.variable_scope('init'):
      x = self._conv('init_conv', x, 3, x.shape[3], 16,
                     self._stride_arr(1))

    strides = [1, 2, 2]
    activate_before_residual = [True, False, False]
    if self.hps.use_bottleneck:
      res_func = self._bottleneck_residual
      filters = [16, 64, 128, 256]
    else:
      res_func = self._residual
      filters = [16, 16, 32, 64]
      # Uncomment the following codes to use w28-10 wide residual
      # network.
      # It is more memory efficient than very deep residual network and
      # has
      # comparably good performance.
      # https://arxiv.org/pdf/1605.07146v1.pdf
      # filters = [16, 160, 320, 640]
      # Update hps.num_residual_units to 4

    with tf.variable_scope('unit_1_0'):
      x = res_func(x, filters[0], filters[1],
                   self._stride_arr(strides[0]),
                   activate_before_residual[0])
    for i in six.moves.range(1, self.hps.num_residual_units):
      with tf.variable_scope('unit_1_%d' % i):
        x = res_func(x, filters[1], filters[1],
                     self._stride_arr(1), False)

    with tf.variable_scope('unit_2_0'):
      x = res_func(x, filters[1], filters[2],
                   self._stride_arr(strides[1]),
                   activate_before_residual[1])
    for i in six.moves.range(1, self.hps.num_residual_units):
      with tf.variable_scope('unit_2_%d' % i):
        x = res_func(x, filters[2], filters[2],
                     self._stride_arr(1), False)

    with tf.variable_scope('unit_3_0'):
      x = res_func(x, filters[2], filters[3],
                   self._stride_arr(strides[2]),
                   activate_before_residual[2])
    for i in six.moves.range(1, self.hps.num_residual_units):
      with tf.variable_scope('unit_3_%d' % i):
        x = res_func(x, filters[3], filters[3],
                     self._stride_arr(1), False)

    with tf.variable_scope('unit_last'):
      x = self._layer_norm('final_bn', x)
      x = self._relu(x, self.hps.relu_leakiness)
      x = self._global_avg_pool(x)

    with tf.variable_scope('logit'):
      logits = self._fully_connected(x, self.hps.nb_classes)
      predictions = tf.nn.softmax(logits)

    return logits, predictions
591925
def build_model(self, n_features, regtype='none'):
        
        self._create_placeholders(n_features)
        self._create_variables(n_features)
        self.encode = self.sample_hidden_from_visible(self.input_data)[0]
        self.reconstruction = self.sample_visible_from_hidden(
            self.encode, n_features)

        hprob0, hstate0, vprob, hprob1, hstate1 = self.gibbs_sampling_step(
            self.input_data, n_features)
        positive = self.compute_positive_association(self.input_data,
                                                     hprob0, hstate0)

        nn_input = vprob

        for step in range(self.gibbs_sampling_steps - 1):
            hprob, hstate, vprob, hprob1, hstate1 = self.gibbs_sampling_step(
                nn_input, n_features)
            nn_input = vprob

        negative = tf.matmul(tf.transpose(vprob), hprob1)

        self.w_upd8 = self.W.assign_add(
            self.learning_rate * (positive - negative) / self.batch_size)

        self.bh_upd8 = self.bh_.assign_add(tf.multiply(self.learning_rate, tf.reduce_mean(
            tf.subtract(hprob0, hprob1), 0)))

        self.bv_upd8 = self.bv_.assign_add(tf.multiply(self.learning_rate, tf.reduce_mean(
            tf.subtract(self.input_data, vprob), 0)))

        variables = [self.W, self.bh_, self.bv_]
        regterm = Layers.regularization(variables, self.regtype, self.regcoef)

        self.cost = self.loss.compile(vprob, self.input_data, regterm=regterm)
1243333
def create_mnist_model(hyper_params, input_shape=(H, W, 1), num_classes=NUM_CLASSES):
    '''
    
    '''
    layers = [
        Conv2D(32, kernel_size=(3, 3), activation='relu', input_shape=input_shape),
        Conv2D(64, (3, 3), activation='relu'),
        MaxPooling2D(pool_size=(2, 2)),
        Flatten(),
        Dense(100, activation='relu'),
        Dense(num_classes, activation='softmax')
    ]

    model = Sequential(layers)

    if hyper_params['optimizer'] == 'Adam':
        optimizer = keras.optimizers.Adam(lr=hyper_params['learning_rate'])
    else:
        optimizer = keras.optimizers.SGD(lr=hyper_params['learning_rate'], momentum=0.9)
    model.compile(loss=keras.losses.categorical_crossentropy, optimizer=optimizer, metrics=['accuracy'])

    return model
263393
def embedding(x,
              vocab_size,
              dense_size,
              name=None,
              reuse=None,
              multiplier=1.0,
              symbol_dropout_rate=0.0,
              embedding_var=None,
              dtype=tf.float32):
  """"""
  with tf.variable_scope(
      name, default_name="embedding", values=[x], reuse=reuse, dtype=dtype):
    if embedding_var is None:
      embedding_var = tf.get_variable("kernel", [vocab_size, dense_size])
    # On the backwards pass, we want to convert the gradient from
    # an indexed-slices to a regular tensor before sending it back to the
    # parameter server. This avoids excess computation on the parameter server.
    if not tf.executing_eagerly():
      embedding_var = convert_gradient_to_tensor(embedding_var)
    x = dropout_no_scaling(x, 1.0 - symbol_dropout_rate)
    emb_x = gather(embedding_var, x, dtype)
    if multiplier != 1.0:
      emb_x *= multiplier
    static_shape = emb_x.shape.as_list()
    if len(static_shape) < 5:
      return emb_x
    assert len(static_shape) == 5
    # If we had an extra channel dimension, assume it's 1, i.e. shape[3] == 1.
    return tf.squeeze(emb_x, 3)
1298642
def _init_evolved_transformer_cache(cache, hparams, batch_size,
                                    attention_init_length, encoder_output,
                                    encoder_decoder_attention_bias,
                                    scope_prefix):
  """"""
  key_channels = hparams.attention_key_channels or hparams.hidden_size
  value_channels = hparams.attention_value_channels or hparams.hidden_size
  num_layers = hparams.num_decoder_layers or hparams.num_hidden_layers
  vars_3d_num_heads = (
      hparams.num_heads if hparams.get("attention_variables_3d") else 0)

  # Add self-attentions.
  if cache is None:
    cache = {}
  cache.update({
      "layer_%d" % layer: {  # pylint: disable=g-complex-comprehension
          _SIXTEEN_HEAD_ATTENTION_NAME: {
              "k":
                  common_attention.split_heads(
                      tf.zeros(
                          [batch_size, attention_init_length, key_channels]),
                      _capped_double_heads(hparams.num_heads)),
              "v":
                  common_attention.split_heads(
                      tf.zeros(
                          [batch_size, attention_init_length, value_channels]),
                      _capped_double_heads(hparams.num_heads)),
          },
          _VANILLA_ATTENTION_NAME: {
              "k":
                  common_attention.split_heads(
                      tf.zeros(
                          [batch_size, attention_init_length, key_channels]),
                      hparams.num_heads),
              "v":
                  common_attention.split_heads(
                      tf.zeros(
                          [batch_size, attention_init_length, value_channels]),
                      hparams.num_heads),
          }
      } for layer in range(num_layers)
  })

  # Add branched layers. Pad with additional zeros for causal convolution.
  for layer in range(num_layers):
    cache["layer_%d" % layer][_CONV_BRANCHES_FIRST_LAYER_NAME] = tf.zeros([
        batch_size, attention_init_length + _DECODER_LEFT_CONV_PADDING,
        hparams.hidden_size
    ])
    cache["layer_%d" % layer][_CONV_BRANCHES_SECOND_LAYER_NAME] = tf.zeros([
        batch_size, attention_init_length + _DECODER_FINAL_CONV_PADDING,
        hparams.hidden_size * 2
    ])

  # Add encoder embedding attentions.
  if encoder_output is not None:
    cache = _add_attend_to_encoder_cache(
        cache=cache,
        attention_name=_FIRST_ATTEND_TO_ENCODER_NAME,
        hparams=hparams,
        num_layers=num_layers,
        key_channels=key_channels,
        value_channels=value_channels,
        vars_3d_num_heads=vars_3d_num_heads,
        scope_prefix=scope_prefix,
        encoder_output=encoder_output)
    cache = _add_attend_to_encoder_cache(
        cache=cache,
        attention_name=_SECOND_ATTEND_TO_ENCODER_NAME,
        hparams=hparams,
        num_layers=num_layers,
        key_channels=key_channels,
        value_channels=value_channels,
        vars_3d_num_heads=vars_3d_num_heads,
        scope_prefix=scope_prefix,
        encoder_output=encoder_output)

    cache["encoder_output"] = encoder_output
    cache["encoder_decoder_attention_bias"] = encoder_decoder_attention_bias

  return cache
86206
def factorization_machine_model(factor_size, num_features,
                                lr_mult_config, wd_mult_config, init_config):
    
    x = mx.symbol.Variable("data", stype='csr')
    # factor, linear and bias terms
    v = mx.symbol.Variable("v", shape=(num_features, factor_size), stype='row_sparse',
                           init=init_config['v'], lr_mult=lr_mult_config['v'],
                           wd_mult=wd_mult_config['v'])
    w = mx.symbol.Variable('w', shape=(num_features, 1), stype='row_sparse',
                      init=init_config['w'], lr_mult=lr_mult_config['w'],
                      wd_mult=wd_mult_config['w'])
    w0 = mx.symbol.Variable('w0', shape=(1,), init=init_config['w0'],
                       lr_mult=lr_mult_config['w0'], wd_mult=wd_mult_config['w0'])
    w1 = mx.symbol.broadcast_add(mx.symbol.dot(x, w), w0)

    # squared terms for subtracting self interactions
    v_s = mx.symbol._internal._square_sum(data=v, axis=1, keepdims=True)
    x_s = x.square()
    bd_sum = mx.sym.dot(x_s, v_s)

    # interactions
    w2 = mx.symbol.dot(x, v)
    w2_squared = 0.5 * mx.symbol.square(data=w2)

    # putting everything together
    w_all = mx.symbol.Concat(w1, w2_squared, dim=1)
    sum1 = w_all.sum(axis=1, keepdims=True)
    sum2 = -0.5 * bd_sum
    model = sum1 + sum2

    y = mx.symbol.Variable("softmax_label")
    model = mx.symbol.LogisticRegressionOutput(data=model, label=y)
    return model
411195
def td_conv(inputs,
            filters,
            kernel_size,
            targeting_count,
            targeting_fn,
            keep_prob,
            is_training,
            do_prune=True,
            strides=(1, 1),
            padding="valid",
            data_format="channels_last",
            dilation_rate=(1, 1),
            activation=None,
            use_bias=True,
            kernel_initializer=None,
            bias_initializer=tf.zeros_initializer(),
            name=None,
            reuse=None):
  """"""
  with tf.variable_scope(name, default_name="td_conv", reuse=reuse):
    nhwc = data_format == "channels_last"
    in_dim = shape_list(inputs)[-1] if nhwc else shape_list(inputs)[1]

    kernel_shape = [kernel_size, kernel_size, in_dim, filters]
    w = tf.get_variable(
        "DW", shape=kernel_shape, initializer=kernel_initializer)
    if use_bias:
      b = tf.get_variable("b", shape=[filters], initializer=bias_initializer)

    if keep_prob < 1.0:
      w = targeted_dropout(
          w,
          targeting_count,
          keep_prob,
          targeting_fn,
          is_training,
          do_prune=do_prune)

    if isinstance(strides, int):
      strides = [strides, strides]
    if isinstance(dilation_rate, int):
      dilation_rate = [dilation_rate, dilation_rate]

    if nhwc:
      strides = [1, strides[0], strides[1], 1]
      dilation_rate = [1, dilation_rate[0], dilation_rate[1], 1]
    else:
      strides = [1, 1, strides[0], strides[1]]
      dilation_rate = [1, 1, dilation_rate[0], dilation_rate[1]]

    y = tf.nn.conv2d(
        inputs,
        w,
        strides,
        padding,
        data_format="NHWC" if nhwc else "NCHW",
        dilations=dilation_rate,
        name=None)

    if use_bias:
      y += b

    if activation:
      y = activation(y)

    return y
668061
def embed(self, x, nfeatures=2):
        
        list_split = 0
        if isinstance(x, list):
            list_split = len(x)
            x = tf.concat(x, 0)

        # pre-process MNIST dataflow data
        x = tf.expand_dims(x, 3)
        x = x * 2 - 1

        # the embedding network
        net = slim.layers.conv2d(x, 20, 5, scope='conv1')
        net = slim.layers.max_pool2d(net, 2, scope='pool1')
        net = slim.layers.conv2d(net, 50, 5, scope='conv2')
        net = slim.layers.max_pool2d(net, 2, scope='pool2')
        net = slim.layers.flatten(net, scope='flatten3')
        net = slim.layers.fully_connected(net, 500, scope='fully_connected4')
        embeddings = slim.layers.fully_connected(net, nfeatures, activation_fn=None, scope='fully_connected5')

        # if "x" was a list of tensors, then split the embeddings
        if list_split > 0:
            embeddings = tf.split(embeddings, list_split, 0)

        return embeddings
325696
def deeper_conv_block(conv_layer, kernel_size, weighted=True):
    '''
    '''
    n_dim = get_n_dim(conv_layer)
    filter_shape = (kernel_size,) * 2
    n_filters = conv_layer.filters
    weight = np.zeros((n_filters, n_filters) + filter_shape)
    center = tuple(map(lambda x: int((x - 1) / 2), filter_shape))
    for i in range(n_filters):
        filter_weight = np.zeros((n_filters,) + filter_shape)
        index = (i,) + center
        filter_weight[index] = 1
        weight[i, ...] = filter_weight
    bias = np.zeros(n_filters)
    new_conv_layer = get_conv_class(n_dim)(
        conv_layer.filters, n_filters, kernel_size=kernel_size
    )
    bn = get_batch_norm_class(n_dim)(n_filters)

    if weighted:
        new_conv_layer.set_weights(
            (add_noise(weight, np.array([0, 1])), add_noise(bias, np.array([0, 1])))
        )
        new_weights = [
            add_noise(np.ones(n_filters, dtype=np.float32), np.array([0, 1])),
            add_noise(np.zeros(n_filters, dtype=np.float32), np.array([0, 1])),
            add_noise(np.zeros(n_filters, dtype=np.float32), np.array([0, 1])),
            add_noise(np.ones(n_filters, dtype=np.float32), np.array([0, 1])),
        ]
        bn.set_weights(new_weights)

    return [StubReLU(), new_conv_layer, bn]
1184594
def lstm(num_hidden, indata, prev_state, param, seqidx, layeridx, dropout=0., num_hidden_proj=0, is_batchnorm=False,
         gamma=None, beta=None, name=None):
    """"""
    # dropout input
    if dropout > 0.:
        indata = mx.sym.Dropout(data=indata, p=dropout)

    i2h = mx.sym.FullyConnected(data=indata,
                                weight=param.i2h_weight,
                                bias=param.i2h_bias,
                                num_hidden=num_hidden * 4,
                                name="t%d_l%d_i2h" % (seqidx, layeridx))
    if is_batchnorm:
        if name is not None:
            i2h = batchnorm(net=i2h, gamma=gamma, beta=beta, name="%s_batchnorm" % name)
        else:
            i2h = batchnorm(net=i2h, gamma=gamma, beta=beta)

    h2h = mx.sym.FullyConnected(data=prev_state.h,
                                weight=param.h2h_weight,
                                # bias=param.h2h_bias,
                                no_bias=True,
                                num_hidden=num_hidden * 4,
                                name="t%d_l%d_h2h" % (seqidx, layeridx))

    gates = i2h + h2h
    slice_gates = mx.sym.SliceChannel(gates, num_outputs=4,
                                      name="t%d_l%d_slice" % (seqidx, layeridx))

    Wcidc = mx.sym.broadcast_mul(param.c2i_bias, prev_state.c) + slice_gates[0]
    in_gate = mx.sym.Activation(Wcidc, act_type="sigmoid")

    in_transform = mx.sym.Activation(slice_gates[1], act_type="tanh")

    Wcfdc = mx.sym.broadcast_mul(param.c2f_bias, prev_state.c) + slice_gates[2]
    forget_gate = mx.sym.Activation(Wcfdc, act_type="sigmoid")

    next_c = (forget_gate * prev_state.c) + (in_gate * in_transform)

    Wcoct = mx.sym.broadcast_mul(param.c2o_bias, next_c) + slice_gates[3]
    out_gate = mx.sym.Activation(Wcoct, act_type="sigmoid")

    next_h = out_gate * mx.sym.Activation(next_c, act_type="tanh")

    if num_hidden_proj > 0:
        proj_next_h = mx.sym.FullyConnected(data=next_h,
                                            weight=param.ph2h_weight,
                                            no_bias=True,
                                            num_hidden=num_hidden_proj,
                                            name="t%d_l%d_ph2h" % (seqidx, layeridx))

        return LSTMState(c=next_c, h=proj_next_h)
    else:
        return LSTMState(c=next_c, h=next_h)
1040257
def next_frame_basic_stochastic_discrete():
  """"""
  hparams = basic_deterministic_params.next_frame_sampling()
  hparams.batch_size = 4
  hparams.video_num_target_frames = 6
  hparams.scheduled_sampling_mode = "prob_inverse_lin"
  hparams.scheduled_sampling_decay_steps = 40000
  hparams.scheduled_sampling_max_prob = 1.0
  hparams.dropout = 0.15
  hparams.filter_double_steps = 3
  hparams.hidden_size = 96
  hparams.learning_rate_constant = 0.002
  hparams.learning_rate_warmup_steps = 2000
  hparams.learning_rate_schedule = "linear_warmup * constant"
  hparams.concat_internal_states = True
  hparams.video_modality_loss_cutoff = 0.03
  hparams.add_hparam("bottleneck_bits", 128)
  hparams.add_hparam("bottleneck_noise", 0.1)
  hparams.add_hparam("discretize_warmup_steps", 40000)
  hparams.add_hparam("latent_rnn_warmup_steps", 40000)
  hparams.add_hparam("latent_rnn_max_sampling", 0.5)
  hparams.add_hparam("latent_use_max_probability", 0.8)
  hparams.add_hparam("full_latent_tower", False)
  hparams.add_hparam("latent_predictor_state_size", 128)
  hparams.add_hparam("latent_predictor_temperature", 1.0)
  hparams.add_hparam("complex_addn", True)
  hparams.add_hparam("recurrent_state_size", 64)
  return hparams
362213
def main( gpu:Param("GPU to run on", str)=None ):
    """"""
    path = Path('/mnt/fe2_disk/')
    tot_epochs,size,bs,lr = 60,224,256,3e-1
    dirname = 'imagenet'

    gpu = setup_distrib(gpu)
    if gpu is None: bs *= torch.cuda.device_count()
    n_gpus = num_distrib() or 1
    workers = min(12, num_cpus()//n_gpus)
    data = get_data(path/dirname, size, bs, workers)
    b_its = len(data.train_dl)//n_gpus

    # Using bs 256 on single GPU as baseline, scale the LR linearly
    tot_bs = bs*n_gpus
    bs_rat = tot_bs/256
    lr *= bs_rat

    ph1 = (TrainingPhase(tot_epochs*0.10*b_its)
            .schedule_hp('lr', (lr/10,lr),  anneal=annealing_cos))
    ph2 = (TrainingPhase(tot_epochs*0.90*b_its)
            .schedule_hp('lr', (lr,lr/1e5), anneal=annealing_cos))
    opt_func = partial(optim.Adam, eps=0.1, betas=(0.9,0.99))
    learn = Learner(data, models.xresnet50(), metrics=[accuracy,top_k_accuracy], wd=1e-3,
        opt_func=opt_func, bn_wd=False, true_wd=True,
        loss_func = LabelSmoothingCrossEntropy()).mixup(alpha=0.2)

    learn.callback_fns += [
        partial(GeneralScheduler, phases=(ph1,ph2)),
        partial(SaveModelCallback, every='epoch', name='model')
    ]
    learn.split(lambda m: (children(m)[-2],))
    if gpu is None: learn.model = nn.DataParallel(learn.model)
    else:           learn.to_distributed(gpu)
    learn.to_fp16(dynamic=True)

    learn.fit(tot_epochs, 1)
    if rank_distrib(): time.sleep(1)
    learn.save('done')
472157
public static int cudnnConvolutionForward(
        cudnnHandle handle, 
        Pointer alpha, 
        cudnnTensorDescriptor xDesc, 
        Pointer x, 
        cudnnFilterDescriptor wDesc, 
        Pointer w, 
        cudnnConvolutionDescriptor convDesc, 
        int algo, 
        Pointer workSpace, 
        long workSpaceSizeInBytes, 
        Pointer beta, 
        cudnnTensorDescriptor yDesc, 
        Pointer y)
    {
        return checkResult(cudnnConvolutionForwardNative(handle, alpha, xDesc, x, wDesc, w, convDesc, algo, workSpace, workSpaceSizeInBytes, beta, yDesc, y));
    }
1245954
def conv2d(ni:int, nf:int, ks:int=3, stride:int=1, padding:int=None, bias=False, init:LayerFunc=nn.init.kaiming_normal_) -> nn.Conv2d:
    ""
    if padding is None: padding = ks//2
    return init_default(nn.Conv2d(ni, nf, kernel_size=ks, stride=stride, padding=padding, bias=bias), init)