new_contents: '''Train'''
  from __future__ import (absolute_import, division, print_function,
                          unicode_literals)
  
  import argparse
  import os
  import sys
  
  import chainer
  import chainer.functions as F
  import chainer.training
  import chainer.training.extensions as extensions
  import numpy as np
  import scipy.io.wavfile as wavfile
  import tqdm
  
  import wavenet.models as models
  import wavenet.utils as utils
  
  
  def generate_and_save_samples(sample_fn, length, count, dir, rate):
      def save_samples(data):
          value = np.iinfo(np.int32).max
          audio = (utils.inverse_mulaw(data) * value).astype(np.int16)
          for idx, sample in enumerate(audio):
              filename = os.path.join(dir, 'sample_{}.wav'.format(idx))
              wavfile.write(filename, rate, sample)
  
      samples = chainer.Variable(
          chainer.cuda.cupy.zeros([count, 1, 1, length], dtype='float32'))
  
      with tqdm.tqdm(total=length) as bar:
          for i in range(length):
              probs = F.softmax(sample_fn(samples))[:, :, 0, 0, i]
              _, level_count = probs.shape
              samples.data[:, 0, 0, i] = chainer.cuda.to_gpu(utils.sample_from(probs.data.get())) / (level_count - 1)
              bar.update()
  
      samples.to_cpu()
      save_samples(samples.data)
  
  def main():
      parser = argparse.ArgumentParser(description='PixelCNN')
      parser.add_argument('--gpu', '-g', type=int, default=-1,
                          help='GPU ID (negative value indicates CPU)')
      parser.add_argument('--model', '-m', default='',
                          help='Path to model for generation')
      parser.add_argument('--hidden_dim', type=int, default=32,
                          help='Number of hidden dimensions')
      parser.add_argument('--out_hidden_dim', type=int, default=32,
                          help='Number of hidden dimensions')
      parser.add_argument('--stacks_num', '-s', type=int, default=5,
                          help='Number of stacks')
      parser.add_argument('--layers_num', '-l', type=int, default=10,
                          help='Number of layers per stack')
      parser.add_argument('--levels', type=int, default=256,
                          help='Level number to quantisize pixel values')
      parser.add_argument('--output', '-o', type=str, default='samples/',
                          help='Output sample directory')
      parser.add_argument('--label', type=np.int32, default=0,
                          help='Class label to generate')
      parser.add_argument('--count', '-c', type=int, default=10,
                          help='Number of samples to generate')
      parser.add_argument('--rate', type=int, default=8000,
                          help='Samples rate')
      parser.add_argument('--length', type=int, default=4096, help='Output sample length')
      args = parser.parse_args()
  
      model = models.WaveNet(args.levels, args.hidden_dim, args.out_hidden_dim, args.stacks_num,
                         args.layers_num, 2)
      if args.gpu >= 0:
          chainer.cuda.get_device(args.gpu).use()
          model.to_gpu()
      chainer.serializers.load_npz(args.model, model)
  
      def sample_fn(samples):
          B, C, H, W = samples.shape
          return model(samples, np.ones(B).astype('i') * args.label)
  
      generate_and_save_samples(sample_fn, args.length, args.count, args.output, args.rate)
  
  
  if __name__ == '__main__':
      sys.exit(main())
  
subject: Add wavenet inference