From 21b08de200d7cfc35d4a32878cf75562470eaaec Mon Sep 17 00:00:00 2001
From: Raindrop <1710404@mail.nankai.edu.cn>
Date: Mon, 21 Mar 2022 15:42:58 +0800
Subject: [PATCH] first jasper commit

---
 PyTorch/contrib/audio/jasper/.keep            |    0
 .../Jasper_pytorch_wuqingdian01/.dockerignore |    8 +
 .../Jasper_pytorch_wuqingdian01/.gitignore    |    9 +
 .../jasper/Jasper_pytorch_wuqingdian01/.keep  |    0
 .../Jasper_pytorch_wuqingdian01/Dockerfile    |   30 +
 .../Jasper_pytorch_wuqingdian01/LICENSE       |  203 ++
 .../jasper/Jasper_pytorch_wuqingdian01/NOTICE |    5 +
 .../Jasper_pytorch_wuqingdian01/README.md     |   49 +
 .../Jasper_pytorch_wuqingdian01/README_raw.md |  843 ++++++
 .../apex/.gitignore                           |  147 +
 .../apex/.gitmodules                          |    7 +
 .../apex/.nojekyll                            |    0
 .../Jasper_pytorch_wuqingdian01/apex/LICENSE  |   11 +
 .../apex/README.md                            |  146 +
 .../apex/apex/RNN/README.md                   |    1 +
 .../apex/apex/RNN/RNNBackend.py               |  365 +++
 .../apex/apex/RNN/__init__.py                 |    3 +
 .../apex/apex/RNN/cells.py                    |   84 +
 .../apex/apex/RNN/models.py                   |   54 +
 .../apex/apex/__init__.py                     |   20 +
 .../apex/apex/_autocast_utils.py              |   17 +
 .../apex/apex/amp/README.md                   |   72 +
 .../apex/apex/amp/__init__.py                 |    5 +
 .../apex/apex/amp/__version__.py              |    2 +
 .../apex/apex/amp/_amp_state.py               |   69 +
 .../apex/apex/amp/_initialize.py              |  263 ++
 .../apex/apex/amp/_process_optimizer.py       |  489 +++
 .../apex/apex/amp/amp.py                      |  177 ++
 .../apex/apex/amp/compat.py                   |   46 +
 .../apex/apex/amp/frontend.py                 |  442 +++
 .../apex/apex/amp/handle.py                   |  281 ++
 .../apex/apex/amp/lists/__init__.py           |    0
 .../apex/amp/lists/functional_overrides.py    |   80 +
 .../apex/apex/amp/lists/tensor_overrides.py   |   63 +
 .../apex/apex/amp/lists/torch_overrides.py    |  115 +
 .../apex/apex/amp/opt.py                      |  103 +
 .../apex/apex/amp/rnn_compat.py               |   53 +
 .../apex/apex/amp/scaler.py                   |  217 ++
 .../apex/apex/amp/utils.py                    |  210 ++
 .../apex/apex/amp/wrap.py                     |  276 ++
 .../apex/apex/contrib/__init__.py             |    0
 .../apex/apex/contrib/bottleneck/__init__.py  |    1 +
 .../apex/contrib/bottleneck/bottleneck.py     |  512 ++++
 .../bottleneck/bottleneck_module_test.py      |  272 ++
 .../apex/apex/contrib/bottleneck/test.py      |   71 +
 .../contrib/csrc/bottleneck/bottleneck.cpp    | 2486 +++++++++++++++
 .../apex/apex/contrib/csrc/fmha/fmha_api.cpp  |  432 +++
 .../apex/apex/contrib/csrc/fmha/src/fmha.h    |  125 +
 .../apex/contrib/csrc/fmha/src/fmha/gemm.h    |  317 ++
 .../contrib/csrc/fmha/src/fmha/gmem_tile.h    |  428 +++
 .../csrc/fmha/src/fmha/kernel_traits.h        |   95 +
 .../apex/contrib/csrc/fmha/src/fmha/mask.h    |   76 +
 .../contrib/csrc/fmha/src/fmha/smem_tile.h    | 1288 ++++++++
 .../apex/contrib/csrc/fmha/src/fmha/softmax.h |  478 +++
 .../apex/contrib/csrc/fmha/src/fmha/utils.h   |  953 ++++++
 .../src/fmha_dgrad_fp16_128_64_kernel.sm80.cu |   60 +
 .../src/fmha_dgrad_fp16_256_64_kernel.sm80.cu |   60 +
 .../src/fmha_dgrad_fp16_384_64_kernel.sm80.cu |   60 +
 .../src/fmha_dgrad_fp16_512_64_kernel.sm80.cu |  105 +
 .../fmha/src/fmha_dgrad_kernel_1xN_reload.h   |  558 ++++
 .../src/fmha_dgrad_kernel_1xN_reload_nl.h     |  571 ++++
 .../src/fmha_fprop_fp16_128_64_kernel.sm80.cu |   58 +
 .../src/fmha_fprop_fp16_256_64_kernel.sm80.cu |   58 +
 .../src/fmha_fprop_fp16_384_64_kernel.sm80.cu |   57 +
 .../src/fmha_fprop_fp16_512_64_kernel.sm80.cu |   98 +
 .../csrc/fmha/src/fmha_fprop_kernel_1xN.h     |  336 +++
 .../csrc/fmha/src/fmha_fprop_kernel_1xN_nl.h  |  343 +++
 .../fmha/src/fmha_fprop_kernel_1xN_reload_v.h |  322 ++
 .../apex/contrib/csrc/fmha/src/fmha_kernel.h  |  169 ++
 .../csrc/fmha/src/fmha_noloop_reduce.cu       |  177 ++
 .../apex/contrib/csrc/fmha/src/fmha_utils.h   |   92 +
 .../apex/contrib/csrc/groupbn/batch_norm.cu   |  328 ++
 .../apex/contrib/csrc/groupbn/batch_norm.h    |  735 +++++
 .../csrc/groupbn/batch_norm_add_relu.cu       |  340 +++
 .../csrc/groupbn/batch_norm_add_relu.h        |  682 +++++
 .../apex/contrib/csrc/groupbn/cuda_utils.h    |   20 +
 .../apex/contrib/csrc/groupbn/interface.cpp   |  175 ++
 .../apex/apex/contrib/csrc/groupbn/ipc.cu     |  129 +
 .../csrc/groupbn/nhwc_batch_norm_kernel.h     | 2685 +++++++++++++++++
 .../apex/apex/contrib/csrc/layer_norm/ln.h    |  200 ++
 .../apex/contrib/csrc/layer_norm/ln_api.cpp   |  246 ++
 .../csrc/layer_norm/ln_bwd_kernels.cuh        |  315 ++
 .../layer_norm/ln_bwd_semi_cuda_kernel.cu     |  234 ++
 .../csrc/layer_norm/ln_fwd_cuda_kernel.cu     |  222 ++
 .../csrc/layer_norm/ln_fwd_kernels.cuh        |  110 +
 .../csrc/layer_norm/ln_kernel_traits.h        |  159 +
 .../apex/contrib/csrc/layer_norm/ln_utils.cuh |  733 +++++
 .../additive_masked_softmax_dropout.cpp       |   91 +
 .../additive_masked_softmax_dropout_cuda.cu   |  131 +
 .../contrib/csrc/multihead_attn/dropout.h     |  306 ++
 .../multihead_attn/encdec_multihead_attn.cpp  |  156 +
 .../encdec_multihead_attn_cuda.cu             |  556 ++++
 .../encdec_multihead_attn_norm_add.cpp        |  198 ++
 .../encdec_multihead_attn_norm_add_cuda.cu    |  640 ++++
 .../contrib/csrc/multihead_attn/layer_norm.h  |  740 +++++
 .../multihead_attn/masked_softmax_dropout.cpp |   93 +
 .../masked_softmax_dropout_cuda.cu            |  147 +
 .../apex/contrib/csrc/multihead_attn/philox.h |   90 +
 .../multihead_attn/self_multihead_attn.cpp    |  132 +
 .../self_multihead_attn_bias.cpp              |  139 +
 ...self_multihead_attn_bias_additive_mask.cpp |  143 +
 ..._multihead_attn_bias_additive_mask_cuda.cu |  463 +++
 .../self_multihead_attn_bias_cuda.cu          |  471 +++
 .../self_multihead_attn_cuda.cu               |  469 +++
 .../self_multihead_attn_norm_add.cpp          |  173 ++
 .../self_multihead_attn_norm_add_cuda.cu      |  556 ++++
 .../contrib/csrc/multihead_attn/softmax.h     | 2641 ++++++++++++++++
 .../multihead_attn/strided_batched_gemm.h     |  407 +++
 .../csrc/optimizers/fused_adam_cuda.cpp       |   86 +
 .../csrc/optimizers/fused_adam_cuda_kernel.cu | 1037 +++++++
 .../csrc/optimizers/fused_lamb_cuda.cpp       |   21 +
 .../csrc/optimizers/fused_lamb_cuda_kernel.cu |  294 ++
 .../optimizers/multi_tensor_distopt_adam.cpp  |   20 +
 .../multi_tensor_distopt_adam_kernel.cu       |  228 ++
 .../optimizers/multi_tensor_distopt_lamb.cpp  |   36 +
 .../multi_tensor_distopt_lamb_kernel.cu       |  506 ++++
 .../csrc/transducer/transducer_joint.cpp      |   98 +
 .../transducer/transducer_joint_kernel.cu     |  973 ++++++
 .../csrc/transducer/transducer_loss.cpp       |  109 +
 .../csrc/transducer/transducer_loss_kernel.cu |  767 +++++
 .../apex/contrib/csrc/xentropy/interface.cpp  |   52 +
 .../contrib/csrc/xentropy/xentropy_kernel.cu  |  718 +++++
 .../func_test_multihead_attn.py               |  108 +
 .../perf_test_multihead_attn.py               |  115 +
 .../apex/apex/contrib/fmha/__init__.py        |    1 +
 .../apex/apex/contrib/fmha/fmha.py            |   74 +
 .../apex/apex/contrib/groupbn/__init__.py     |    9 +
 .../apex/apex/contrib/groupbn/batch_norm.py   |  225 ++
 .../apex/apex/contrib/layer_norm/__init__.py  |    1 +
 .../apex/contrib/layer_norm/layer_norm.py     |   53 +
 .../apex/contrib/multihead_attn/README.md     |   60 +
 .../apex/contrib/multihead_attn/__init__.py   |    3 +
 .../multihead_attn/encdec_multihead_attn.py   |  141 +
 .../encdec_multihead_attn_func.py             |  268 ++
 .../fast_encdec_multihead_attn_func.py        |   88 +
 ...ast_encdec_multihead_attn_norm_add_func.py |  130 +
 .../fast_self_multihead_attn_func.py          |  196 ++
 .../fast_self_multihead_attn_norm_add_func.py |  106 +
 .../mask_softmax_dropout_func.py              |   81 +
 .../multihead_attn/self_multihead_attn.py     |  178 ++
 .../self_multihead_attn_func.py               |  235 ++
 .../apex/apex/contrib/optimizers/__init__.py  |    3 +
 .../optimizers/distributed_fused_adam.py      |  636 ++++
 .../optimizers/distributed_fused_adam_v2.py   |  615 ++++
 .../optimizers/distributed_fused_adam_v3.py   |  325 ++
 .../optimizers/distributed_fused_lamb.py      |  975 ++++++
 .../apex/contrib/optimizers/fp16_optimizer.py |  243 ++
 .../apex/contrib/optimizers/fused_adam.py     |  206 ++
 .../apex/contrib/optimizers/fused_lamb.py     |  208 ++
 .../apex/apex/contrib/optimizers/fused_sgd.py |  211 ++
 .../apex/apex/contrib/sparsity/README.md      |   78 +
 .../apex/apex/contrib/sparsity/__init__.py    |    2 +
 .../apex/apex/contrib/sparsity/asp.py         |  217 ++
 .../apex/contrib/sparsity/sparse_masklib.py   |  184 ++
 .../sparsity/test/checkpointing_test_part1.py |   94 +
 .../sparsity/test/checkpointing_test_part2.py |   79 +
 .../test/checkpointing_test_reference.py      |   96 +
 .../apex/contrib/sparsity/test/toy_problem.py |   87 +
 .../apex/apex/contrib/test/fmha/test_fmha.py  |  121 +
 .../test/fused_dense/test_fused_dense.py      |   44 +
 .../test/layer_norm/test_fast_layer_norm.py   |  275 ++
 .../test_encdec_multihead_attn.py             |  136 +
 .../test_encdec_multihead_attn_norm_add.py    |   77 +
 .../test_fast_self_multihead_attn_bias.py     |   77 +
 .../multihead_attn/test_mha_fused_softmax.py  |   42 +
 .../test_self_multihead_attn.py               |  130 +
 .../test_self_multihead_attn_norm_add.py      |   72 +
 .../apex/contrib/test/test_label_smoothing.py |  128 +
 .../test/transducer/test_transducer_joint.py  |  157 +
 .../test/transducer/test_transducer_loss.py   |  133 +
 .../contrib/test/transducer/transducer_ref.py |  112 +
 .../apex/apex/contrib/transducer/__init__.py  |    2 +
 .../apex/contrib/transducer/transducer.py     |  195 ++
 .../apex/apex/contrib/xentropy/__init__.py    |    9 +
 .../apex/contrib/xentropy/softmax_xentropy.py |   28 +
 .../apex/apex/fp16_utils/README.md            |   16 +
 .../apex/apex/fp16_utils/__init__.py          |   16 +
 .../apex/apex/fp16_utils/fp16_optimizer.py    |  554 ++++
 .../apex/apex/fp16_utils/fp16util.py          |  187 ++
 .../apex/apex/fp16_utils/loss_scaler.py       |  186 ++
 .../apex/apex/fused_dense/__init__.py         |    1 +
 .../apex/apex/fused_dense/fused_dense.py      |   85 +
 .../apex/apex/mlp/__init__.py                 |    1 +
 .../apex/apex/mlp/mlp.py                      |   79 +
 .../apex/apex/multi_tensor_apply/__init__.py  |    4 +
 .../multi_tensor_apply/multi_tensor_apply.py  |   30 +
 .../apex/apex/normalization/__init__.py       |    1 +
 .../apex/normalization/fused_layer_norm.py    |  218 ++
 .../apex/apex/optimizers/__init__.py          |    5 +
 .../apex/apex/optimizers/fused_adagrad.py     |  122 +
 .../apex/apex/optimizers/fused_adam.py        |  173 ++
 .../apex/apex/optimizers/fused_lamb.py        |  215 ++
 .../apex/apex/optimizers/fused_novograd.py    |  214 ++
 .../apex/apex/optimizers/fused_sgd.py         |  227 ++
 .../apex/apex/parallel/LARC.py                |  107 +
 .../apex/apex/parallel/README.md              |   66 +
 .../apex/apex/parallel/__init__.py            |   95 +
 .../apex/apex/parallel/distributed.py         |  639 ++++
 .../apex/apex/parallel/multiproc.py           |   35 +
 .../apex/parallel/optimized_sync_batchnorm.py |   85 +
 .../optimized_sync_batchnorm_kernel.py        |  119 +
 .../apex/apex/parallel/sync_batchnorm.py      |  134 +
 .../apex/parallel/sync_batchnorm_kernel.py    |   87 +
 .../apex/apex/pyprof/FAQs.md                  |   21 +
 .../apex/apex/pyprof/README.md                |  252 ++
 .../apex/apex/pyprof/__init__.py              |    3 +
 .../apex/apex/pyprof/examples/.gitignore      |    4 +
 .../apex/apex/pyprof/examples/apex/README.md  |    1 +
 .../apex/pyprof/examples/apex/fused_adam.py   |   20 +
 .../pyprof/examples/apex/fused_layer_norm.py  |   28 +
 .../apex/apex/pyprof/examples/apex/test.sh    |   30 +
 .../examples/custom_func_module/README.md     |    1 +
 .../custom_func_module/custom_function.py     |   33 +
 .../custom_func_module/custom_module.py       |   27 +
 .../examples/custom_func_module/test.sh       |   30 +
 .../apex/pyprof/examples/imagenet/imagenet.py |  135 +
 .../apex/pyprof/examples/imagenet/test.sh     |   36 +
 .../apex/apex/pyprof/examples/jit/README.md   |   14 +
 .../examples/jit/jit_script_function.py       |   30 +
 .../pyprof/examples/jit/jit_script_method.py  |   31 +
 .../pyprof/examples/jit/jit_trace_function.py |   30 +
 .../pyprof/examples/jit/jit_trace_method.py   |   36 +
 .../apex/apex/pyprof/examples/jit/test.sh     |   30 +
 .../apex/apex/pyprof/examples/lenet.py        |   65 +
 .../apex/apex/pyprof/examples/operators.py    |  145 +
 .../apex/apex/pyprof/examples/simple.py       |   38 +
 .../pyprof/examples/user_annotation/README.md |   21 +
 .../pyprof/examples/user_annotation/resnet.py |  215 ++
 .../pyprof/examples/user_annotation/test.sh   |   31 +
 .../apex/apex/pyprof/nvtx/__init__.py         |    2 +
 .../apex/apex/pyprof/nvtx/nvmarker.py         |  222 ++
 .../apex/apex/pyprof/parse/__init__.py        |    0
 .../apex/apex/pyprof/parse/__main__.py        |   10 +
 .../apex/apex/pyprof/parse/db.py              |   61 +
 .../apex/apex/pyprof/parse/kernel.py          |  210 ++
 .../apex/apex/pyprof/parse/nvvp.py            |  282 ++
 .../apex/apex/pyprof/parse/parse.py           |  122 +
 .../apex/apex/pyprof/prof/__init__.py         |    1 +
 .../apex/apex/pyprof/prof/__main__.py         |   10 +
 .../apex/apex/pyprof/prof/activation.py       |   65 +
 .../apex/apex/pyprof/prof/base.py             |   47 +
 .../apex/apex/pyprof/prof/blas.py             |  340 +++
 .../apex/apex/pyprof/prof/conv.py             |  236 ++
 .../apex/apex/pyprof/prof/convert.py          |   62 +
 .../apex/apex/pyprof/prof/data.py             |   54 +
 .../apex/apex/pyprof/prof/dropout.py          |   50 +
 .../apex/apex/pyprof/prof/embedding.py        |   71 +
 .../pyprof/prof/index_slice_join_mutate.py    |  419 +++
 .../apex/apex/pyprof/prof/linear.py           |  188 ++
 .../apex/apex/pyprof/prof/loss.py             |   84 +
 .../apex/apex/pyprof/prof/misc.py             |  219 ++
 .../apex/apex/pyprof/prof/normalization.py    |   54 +
 .../apex/apex/pyprof/prof/optim.py            |   65 +
 .../apex/apex/pyprof/prof/output.py           |  149 +
 .../apex/apex/pyprof/prof/pointwise.py        |  166 +
 .../apex/apex/pyprof/prof/pooling.py          |   59 +
 .../apex/apex/pyprof/prof/prof.py             |  256 ++
 .../apex/apex/pyprof/prof/randomSample.py     |   43 +
 .../apex/apex/pyprof/prof/recurrentCell.py    |  207 ++
 .../apex/apex/pyprof/prof/reduction.py        |  150 +
 .../apex/apex/pyprof/prof/softmax.py          |  115 +
 .../apex/apex/pyprof/prof/usage.py            |   73 +
 .../apex/apex/pyprof/prof/utility.py          |   60 +
 .../apex/apex/reparameterization/README.md    |    1 +
 .../apex/apex/reparameterization/__init__.py  |  127 +
 .../reparameterization/reparameterization.py  |  151 +
 .../apex/reparameterization/weight_norm.py    |   78 +
 .../apex/apex/transformer/README.md           |   81 +
 .../apex/apex/transformer/__init__.py         |   21 +
 .../apex/apex/transformer/enums.py            |   30 +
 .../apex/transformer/functional/__init__.py   |    5 +
 .../transformer/functional/fused_softmax.py   |  199 ++
 .../apex/apex/transformer/microbatches.py     |  172 ++
 .../apex/apex/transformer/parallel_state.py   |  382 +++
 .../transformer/pipeline_parallel/__init__.py |    8 +
 .../transformer/pipeline_parallel/_timers.py  |   83 +
 .../pipeline_parallel/p2p_communication.py    |  398 +++
 .../pipeline_parallel/schedules/__init__.py   |   39 +
 .../pipeline_parallel/schedules/common.py     |  218 ++
 .../schedules/fwd_bwd_no_pipelining.py        |   81 +
 .../fwd_bwd_pipelining_with_interleaving.py   |  298 ++
 ...fwd_bwd_pipelining_without_interleaving.py |  175 ++
 .../transformer/pipeline_parallel/utils.py    |  317 ++
 .../transformer/tensor_parallel/__init__.py   |   74 +
 .../tensor_parallel/cross_entropy.py          |  103 +
 .../apex/transformer/tensor_parallel/data.py  |  113 +
 .../transformer/tensor_parallel/layers.py     |  477 +++
 .../transformer/tensor_parallel/mappings.py   |  159 +
 .../transformer/tensor_parallel/memory.py     |  136 +
 .../transformer/tensor_parallel/random.py     |  294 ++
 .../apex/transformer/tensor_parallel/utils.py |   54 +
 .../apex/apex/transformer/testing/__init__.py |    0
 .../apex/transformer/testing/arguments.py     |  806 +++++
 .../apex/apex/transformer/testing/commons.py  |   87 +
 .../apex/transformer/testing/global_vars.py   |  260 ++
 .../transformer/testing/standalone_gpt.py     | 1506 +++++++++
 .../apex/apex/transformer/utils.py            |   47 +
 .../apex/csrc/amp_C_frontend.cpp              |  145 +
 .../apex/csrc/compat.h                        |    9 +
 .../apex/csrc/flatten_unflatten.cpp           |   18 +
 .../apex/csrc/fused_dense.cpp                 |  192 ++
 .../apex/csrc/fused_dense_cuda.cu             | 1437 +++++++++
 .../apex/csrc/layer_norm_cuda.cpp             |  267 ++
 .../apex/csrc/layer_norm_cuda_kernel.cu       |  830 +++++
 .../csrc/megatron/scaled_masked_softmax.cpp   |   97 +
 .../csrc/megatron/scaled_masked_softmax.h     |  505 ++++
 .../megatron/scaled_masked_softmax_cuda.cu    |  117 +
 .../scaled_upper_triang_masked_softmax.cpp    |   72 +
 .../scaled_upper_triang_masked_softmax.h      |  513 ++++
 ...scaled_upper_triang_masked_softmax_cuda.cu |   98 +
 .../apex/csrc/mlp.cpp                         |  166 +
 .../apex/csrc/mlp_cuda.cu                     | 1678 ++++++++++
 .../apex/csrc/multi_tensor_adagrad.cu         |  100 +
 .../apex/csrc/multi_tensor_adam.cu            |  171 ++
 .../apex/csrc/multi_tensor_apply.cuh          |  133 +
 .../apex/csrc/multi_tensor_axpby_kernel.cu    |  157 +
 .../apex/csrc/multi_tensor_l2norm_kernel.cu   |  448 +++
 .../csrc/multi_tensor_l2norm_scale_kernel.cu  |  326 ++
 .../apex/csrc/multi_tensor_lamb.cu            |  413 +++
 .../apex/csrc/multi_tensor_lamb_stage_1.cu    |  151 +
 .../apex/csrc/multi_tensor_lamb_stage_2.cu    |  125 +
 .../apex/csrc/multi_tensor_novograd.cu        |  188 ++
 .../apex/csrc/multi_tensor_scale_kernel.cu    |  136 +
 .../apex/csrc/multi_tensor_sgd_kernel.cu      |  280 ++
 .../apex/csrc/syncbn.cpp                      |  109 +
 .../apex/csrc/type_shim.h                     |  387 +++
 .../apex/csrc/welford.cu                      | 1510 +++++++++
 .../apex/docs/Makefile                        |   32 +
 .../docs/source/_static/css/pytorch_theme.css |  118 +
 .../apex/docs/source/_templates/layout.html   |   51 +
 .../apex/docs/source/conf.py                  |  248 ++
 .../apex/examples/README.md                   |    4 +
 .../apex/examples/dcgan/README.md             |   41 +
 .../apex/examples/dcgan/main_amp.py           |  274 ++
 .../apex/examples/docker/Dockerfile           |   16 +
 .../apex/examples/docker/README.md            |   40 +
 .../apex/examples/imagenet/README.md          |  183 ++
 .../apex/examples/imagenet/main_amp.py        |  543 ++++
 .../examples/simple/distributed/README.md     |   13 +
 .../distributed/distributed_data_parallel.py  |   65 +
 .../apex/examples/simple/distributed/run.sh   |    2 +
 .../apex/requirements.txt                     |    5 +
 .../apex/requirements_dev.txt                 |    3 +
 .../Jasper_pytorch_wuqingdian01/apex/setup.py |  544 ++++
 .../apex/tests/L0/run_amp/__init__.py         |    0
 .../tests/L0/run_amp/test_add_param_group.py  |  148 +
 .../apex/tests/L0/run_amp/test_basic_casts.py |  143 +
 .../apex/tests/L0/run_amp/test_cache.py       |  137 +
 .../tests/L0/run_amp/test_checkpointing.py    |  267 ++
 .../apex/tests/L0/run_amp/test_fused_sgd.py   |  794 +++++
 .../apex/tests/L0/run_amp/test_larc.py        |   53 +
 .../L0/run_amp/test_multi_tensor_axpby.py     |  180 ++
 .../L0/run_amp/test_multi_tensor_l2norm.py    |   87 +
 .../L0/run_amp/test_multi_tensor_scale.py     |  126 +
 .../test_multiple_models_optimizers_losses.py |  762 +++++
 .../apex/tests/L0/run_amp/test_promotion.py   |   75 +
 .../apex/tests/L0/run_amp/test_rnn.py         |  116 +
 .../apex/tests/L0/run_amp/utils.py            |   21 +
 .../apex/tests/L0/run_fp16util/__init__.py    |    0
 .../tests/L0/run_fp16util/test_fp16util.py    |   75 +
 .../test_fused_layer_norm.py                  |  143 +
 .../apex/tests/L0/run_mlp/test_mlp.py         |  218 ++
 .../apex/tests/L0/run_optimizers/__init__.py  |    0
 .../tests/L0/run_optimizers/test_dist_adam.py |  183 ++
 .../L0/run_optimizers/test_fused_novograd.py  |  170 ++
 .../L0/run_optimizers/test_fused_optimizer.py |  284 ++
 .../apex/tests/L0/run_optimizers/test_lamb.py |  270 ++
 .../apex/tests/L0/run_pyprof_data/__init__.py |    0
 .../L0/run_pyprof_data/test_pyprof_data.py    |   43 +
 .../apex/tests/L0/run_pyprof_nvtx/__init__.py |    1 +
 .../L0/run_pyprof_nvtx/test_pyprof_nvtx.py    |  526 ++++
 .../apex/tests/L0/run_test.py                 |   72 +
 .../apex/tests/L0/run_transformer/__init__.py |    0
 .../run_transformer/run_cross_entropy_test.py |  110 +
 .../tests/L0/run_transformer/run_data_test.py |   95 +
 .../L0/run_transformer/run_initialize_test.py |  104 +
 .../L0/run_transformer/run_layers_test.py     |  696 +++++
 .../L0/run_transformer/run_mappings_test.py   |   63 +
 .../run_megatron_gpt_pipeline.py              |  132 +
 .../run_pipeline_parallel_test.py             |  198 ++
 .../L0/run_transformer/run_random_test.py     |  213 ++
 .../L0/run_transformer/run_utils_test.py      |   22 +
 .../L0/run_transformer/test_batch_sampler.py  |  149 +
 .../L0/run_transformer/test_fused_softmax.py  |  168 ++
 .../test_transformer_module.py                |   89 +
 .../apex/tests/L1/common/compare.py           |   64 +
 .../apex/tests/L1/common/main_amp.py          |  526 ++++
 .../apex/tests/L1/common/run_test.sh          |  144 +
 .../apex/tests/L1/cross_product/run.sh        |    6 +
 .../tests/L1/cross_product_distributed/run.sh |    4 +
 .../DDP/ddp_race_condition_test.py            |   69 +
 .../tests/distributed/DDP/run_race_test.sh    |    3 +
 .../amp_master_params/amp_master_params.py    |   70 +
 .../distributed/amp_master_params/compare.py  |   28 +
 .../distributed/amp_master_params/run.sh      |    4 +
 .../python_single_gpu_unit_test.py            |  111 +
 .../synced_batchnorm/single_gpu_unit_test.py  |  159 +
 .../synced_batchnorm/test_batchnorm1d.py      |   18 +
 .../synced_batchnorm/test_groups.py           |  185 ++
 .../two_gpu_test_different_batch_size.py      |  158 +
 .../synced_batchnorm/two_gpu_unit_test.py     |  180 ++
 .../distributed/synced_batchnorm/unit_test.sh |    8 +
 .../apex/tests/docker_extension_builds/run.sh |   73 +
 .../common/__init__.py                        |    0
 .../common/audio.py                           |  247 ++
 .../common/dali/__init__.py                   |   13 +
 .../common/dali/data_loader.py                |  158 +
 .../common/dali/iterator.py                   |  161 +
 .../common/dali/pipeline.py                   |  366 +++
 .../common/dataset.py                         |  237 ++
 .../common/features.py                        |  522 ++++
 .../common/helpers.py                         |  300 ++
 .../common/metrics.py                         |   59 +
 .../common/optimizers.py                      |  269 ++
 .../common/tb_dllogger.py                     |  159 +
 .../common/text/LICENSE                       |   19 +
 .../common/text/__init__.py                   |   32 +
 .../common/text/cleaners.py                   |  107 +
 .../common/text/numbers.py                    |   99 +
 .../common/text/symbols.py                    |   19 +
 .../common/train.py                           |  516 ++++
 .../common/utils.py                           |   20 +
 .../configs/jasper10x5dr_speca.yaml           |  139 +
 .../configs/jasper10x5dr_speedp-offline.yaml  |  139 +
 .../jasper10x5dr_speedp-offline_speca.yaml    |  139 +
 ...per10x5dr_speedp-offline_speca_nomask.yaml |  139 +
 .../jasper10x5dr_speedp-online-discrete.yaml  |  144 +
 ...er10x5dr_speedp-online-discrete_speca.yaml |  144 +
 .../jasper10x5dr_speedp-online_speca.yaml     |  144 +
 .../Jasper_pytorch_wuqingdian01/inference.py  |  398 +++
 .../jasper/config.py                          |  125 +
 .../jasper/model.py                           |  275 ++
 .../notebooks/README.md                       |  203 ++
 .../platform/DGX1-16GB_Jasper_AMP_8GPU.sh     |    3 +
 .../platform/DGX1-16GB_Jasper_FP32_8GPU.sh    |    3 +
 .../platform/DGX1-32GB_Jasper_AMP_8GPU.sh     |    3 +
 .../platform/DGX1-32GB_Jasper_FP32_8GPU.sh    |    3 +
 .../platform/DGX2_Jasper_AMP_16GPU.sh         |    3 +
 .../platform/DGX2_Jasper_AMP_8GPU.sh          |    3 +
 .../platform/DGX2_Jasper_FP32_16GPU.sh        |    3 +
 .../platform/DGX2_Jasper_FP32_8GPU.sh         |    3 +
 .../platform/DGXA100_Jasper_AMP_8GPU.sh       |    3 +
 .../platform/DGXA100_Jasper_TF32_8GPU.sh      |    3 +
 .../requirements.txt                          |   13 +
 .../scripts/docker/build.sh                   |    3 +
 .../scripts/docker/launch.sh                  |   31 +
 .../scripts/download_librispeech.sh           |   32 +
 .../scripts/evaluation.sh                     |   22 +
 .../scripts/inference.sh                      |   65 +
 .../scripts/inference_benchmark.sh            |   38 +
 .../scripts/preprocess_librispeech.sh         |   54 +
 .../scripts/train.sh                          |   90 +
 .../scripts/train_benchmark.sh                |   49 +
 .../test/docker/build.sh                      |    3 +
 .../test/docker/launch.sh                     |   31 +
 .../test/env_npu.sh                           |   80 +
 .../test/train_eval_8p.sh                     |   65 +
 .../test/train_full_1p.sh                     |   91 +
 .../test/train_full_8p.sh                     |   90 +
 .../test/train_performance_1p.sh              |   95 +
 .../test/train_performance_8p.sh              |   95 +
 .../Jasper_pytorch_wuqingdian01/train.py      |  528 ++++
 .../triton/Dockerfile                         |   10 +
 .../triton/README.md                          |  388 +++
 .../triton/converter.py                       |  254 ++
 .../triton/jasper-client.py                   |  400 +++
 .../triton/jasper_module.py                   |  174 ++
 .../fp16/decoder-ts-script/config.pbtxt       |   46 +
 .../feature-extractor-ts-trace/config.pbtxt   |   32 +
 .../fp16/jasper-onnx-ensemble/config.pbtxt    |   63 +
 .../fp16/jasper-onnx/config.pbtxt             |   58 +
 .../jasper-tensorrt-ensemble/config.pbtxt     |   63 +
 .../fp16/jasper-tensorrt/config.pbtxt         |   58 +
 .../jasper-ts-trace-ensemble/config.pbtxt     |   63 +
 .../fp16/jasper-ts-trace/config.pbtxt         |   32 +
 .../fp32/decoder-ts-script/config.pbtxt       |   46 +
 .../feature-extractor-ts-trace/config.pbtxt   |   32 +
 .../fp32/jasper-onnx-ensemble/config.pbtxt    |   63 +
 .../fp32/jasper-onnx/config.pbtxt             |   58 +
 .../jasper-tensorrt-ensemble/config.pbtxt     |   63 +
 .../fp32/jasper-tensorrt/config.pbtxt         |   58 +
 .../jasper-ts-trace-ensemble/config.pbtxt     |   63 +
 .../fp32/jasper-ts-trace/config.pbtxt         |   32 +
 .../triton/pytorch/__init__.py                |    0
 .../triton/pytorch/utils.py                   |  263 ++
 .../scripts/docker/build_triton_client.sh     |    7 +
 .../scripts/download_triton_librispeech.sh    |   29 +
 .../triton/scripts/execute_all_perf_runs.sh   |   77 +
 .../triton/scripts/export_model.sh            |   78 +
 .../triton/scripts/generate_perf_results.sh   |  108 +
 .../scripts/prepare_model_repository.sh       |   54 +
 .../scripts/preprocess_triton_librispeech.sh  |   21 +
 .../triton/scripts/run_client.sh              |   49 +
 .../triton/scripts/run_perf_client.sh         |   86 +
 .../triton/scripts/run_server.sh              |   72 +
 .../triton/scripts/wait_for_triton_server.sh  |   36 +
 .../triton/speech_utils.py                    |  417 +++
 .../triton/tensorrt_io_props_fp16.json        |    1 +
 .../triton/tensorrt_io_props_fp32.json        |    1 +
 .../triton/triton_librispeech.csv             |    2 +
 .../utils/__init__.py                         |    0
 .../utils/convert_librispeech.py              |   82 +
 .../utils/download_librispeech.py             |   72 +
 .../utils/download_utils.py                   |   71 +
 .../utils/inference_librispeech.csv           |    5 +
 .../utils/librispeech.csv                     |    8 +
 .../utils/preprocessing_utils.py              |   76 +
 507 files changed, 93003 insertions(+)
 create mode 100644 PyTorch/contrib/audio/jasper/.keep
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.dockerignore
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.gitignore
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.keep
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/Dockerfile
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/LICENSE
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/NOTICE
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/README_raw.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.gitignore
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.gitmodules
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.nojekyll
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/LICENSE
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/RNNBackend.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/cells.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/models.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/_autocast_utils.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/__version__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_amp_state.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_initialize.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_process_optimizer.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/amp.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/compat.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/frontend.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/handle.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/functional_overrides.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/tensor_overrides.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/torch_overrides.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/opt.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/rnn_compat.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/scaler.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/utils.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/wrap.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/bottleneck.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/bottleneck_module_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/bottleneck/bottleneck.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/fmha_api.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/gemm.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/gmem_tile.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/kernel_traits.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/mask.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/smem_tile.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/softmax.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/utils.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_128_64_kernel.sm80.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_256_64_kernel.sm80.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_384_64_kernel.sm80.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_512_64_kernel.sm80.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_kernel_1xN_reload.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_kernel_1xN_reload_nl.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_128_64_kernel.sm80.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_256_64_kernel.sm80.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_384_64_kernel.sm80.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_512_64_kernel.sm80.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN_nl.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN_reload_v.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_kernel.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_noloop_reduce.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_utils.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm_add_relu.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm_add_relu.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/cuda_utils.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/interface.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/ipc.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/nhwc_batch_norm_kernel.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_api.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_bwd_kernels.cuh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_bwd_semi_cuda_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_fwd_cuda_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_fwd_kernels.cuh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_kernel_traits.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_utils.cuh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/dropout.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/layer_norm.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/masked_softmax_dropout.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/masked_softmax_dropout_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/philox.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/softmax.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/strided_batched_gemm.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_adam_cuda.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_adam_cuda_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_lamb_cuda.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_lamb_cuda_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_adam.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_adam_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_joint.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_joint_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_loss.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_loss_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/xentropy/interface.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/xentropy/xentropy_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/examples/multihead_attn/func_test_multihead_attn.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/examples/multihead_attn/perf_test_multihead_attn.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/fmha/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/fmha/fmha.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/groupbn/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/groupbn/batch_norm.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/layer_norm/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/layer_norm/layer_norm.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/encdec_multihead_attn.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/encdec_multihead_attn_func.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_encdec_multihead_attn_func.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_encdec_multihead_attn_norm_add_func.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_self_multihead_attn_func.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_self_multihead_attn_norm_add_func.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/mask_softmax_dropout_func.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/self_multihead_attn.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/self_multihead_attn_func.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam_v2.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam_v3.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_lamb.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fp16_optimizer.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_adam.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_lamb.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_sgd.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/asp.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/sparse_masklib.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_part1.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_part2.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_reference.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/toy_problem.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/fmha/test_fmha.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/fused_dense/test_fused_dense.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/layer_norm/test_fast_layer_norm.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_encdec_multihead_attn.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_encdec_multihead_attn_norm_add.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_fast_self_multihead_attn_bias.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_mha_fused_softmax.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_self_multihead_attn.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_self_multihead_attn_norm_add.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/test_label_smoothing.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/test_transducer_joint.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/test_transducer_loss.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/transducer_ref.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/transducer/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/transducer/transducer.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/xentropy/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/xentropy/softmax_xentropy.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/fp16_optimizer.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/fp16util.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/loss_scaler.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fused_dense/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fused_dense/fused_dense.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/mlp/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/mlp/mlp.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/multi_tensor_apply/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/multi_tensor_apply/multi_tensor_apply.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/normalization/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/normalization/fused_layer_norm.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_adagrad.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_adam.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_lamb.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_novograd.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_sgd.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/LARC.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/distributed.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/multiproc.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/optimized_sync_batchnorm.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/optimized_sync_batchnorm_kernel.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/sync_batchnorm.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/sync_batchnorm_kernel.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/FAQs.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/.gitignore
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/fused_adam.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/fused_layer_norm.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/test.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/custom_function.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/custom_module.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/test.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/imagenet/imagenet.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/imagenet/test.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_script_function.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_script_method.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_trace_function.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_trace_method.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/test.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/lenet.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/operators.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/simple.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/resnet.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/test.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/nvtx/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/nvtx/nvmarker.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/__main__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/db.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/kernel.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/nvvp.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/parse.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/__main__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/activation.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/base.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/blas.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/conv.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/convert.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/data.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/dropout.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/embedding.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/index_slice_join_mutate.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/linear.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/loss.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/misc.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/normalization.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/optim.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/output.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/pointwise.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/pooling.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/prof.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/randomSample.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/recurrentCell.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/reduction.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/softmax.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/usage.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/utility.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/reparameterization.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/weight_norm.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/enums.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/functional/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/functional/fused_softmax.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/microbatches.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/parallel_state.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/_timers.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/p2p_communication.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/common.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/utils.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/cross_entropy.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/data.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/layers.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/mappings.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/memory.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/random.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/utils.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/arguments.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/commons.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/global_vars.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/standalone_gpt.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/utils.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/amp_C_frontend.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/compat.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/flatten_unflatten.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/fused_dense.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/fused_dense_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/layer_norm_cuda.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/layer_norm_cuda_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/mlp.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/mlp_cuda.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_adagrad.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_adam.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_apply.cuh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_axpby_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_l2norm_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_l2norm_scale_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb_stage_1.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb_stage_2.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_novograd.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_scale_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_sgd_kernel.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/syncbn.cpp
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/type_shim.h
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/welford.cu
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/Makefile
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/_static/css/pytorch_theme.css
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/_templates/layout.html
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/conf.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/dcgan/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/dcgan/main_amp.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/docker/Dockerfile
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/docker/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/imagenet/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/imagenet/main_amp.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/distributed_data_parallel.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/run.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/requirements.txt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/requirements_dev.txt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/setup.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_add_param_group.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_basic_casts.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_cache.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_checkpointing.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_fused_sgd.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_larc.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_axpby.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_l2norm.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_scale.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multiple_models_optimizers_losses.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_promotion.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_rnn.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/utils.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fp16util/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fp16util/test_fp16util.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_mlp/test_mlp.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_dist_adam.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_fused_novograd.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_fused_optimizer.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_lamb.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_data/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_data/test_pyprof_data.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_nvtx/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_nvtx/test_pyprof_nvtx.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_cross_entropy_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_data_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_initialize_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_layers_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_mappings_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_megatron_gpt_pipeline.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_pipeline_parallel_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_random_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_utils_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_batch_sampler.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_fused_softmax.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_transformer_module.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/compare.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/main_amp.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/run_test.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/cross_product/run.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/cross_product_distributed/run.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/DDP/ddp_race_condition_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/DDP/run_race_test.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/amp_master_params.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/compare.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/run.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/python_single_gpu_unit_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/single_gpu_unit_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/test_batchnorm1d.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/test_groups.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/two_gpu_test_different_batch_size.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/two_gpu_unit_test.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/unit_test.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/docker_extension_builds/run.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/audio.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/data_loader.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/iterator.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/pipeline.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dataset.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/features.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/helpers.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/metrics.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/optimizers.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/tb_dllogger.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/LICENSE
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/cleaners.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/numbers.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/symbols.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/train.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/utils.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speca.yaml
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline.yaml
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline_speca.yaml
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline_speca_nomask.yaml
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online-discrete.yaml
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online-discrete_speca.yaml
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online_speca.yaml
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/inference.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/jasper/config.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/jasper/model.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/notebooks/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-16GB_Jasper_AMP_8GPU.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-16GB_Jasper_FP32_8GPU.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-32GB_Jasper_AMP_8GPU.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-32GB_Jasper_FP32_8GPU.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_AMP_16GPU.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_AMP_8GPU.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_FP32_16GPU.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_FP32_8GPU.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGXA100_Jasper_AMP_8GPU.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGXA100_Jasper_TF32_8GPU.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/requirements.txt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/docker/build.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/docker/launch.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/download_librispeech.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/evaluation.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/inference.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/inference_benchmark.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/preprocess_librispeech.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/train.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/train_benchmark.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/docker/build.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/docker/launch.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/env_npu.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_eval_8p.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_full_1p.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_full_8p.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_performance_1p.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_performance_8p.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/train.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/Dockerfile
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/README.md
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/converter.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/jasper-client.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/jasper_module.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/decoder-ts-script/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/feature-extractor-ts-trace/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-onnx-ensemble/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-onnx/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-tensorrt-ensemble/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-tensorrt/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-ts-trace-ensemble/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-ts-trace/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/decoder-ts-script/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/feature-extractor-ts-trace/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-onnx-ensemble/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-onnx/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-tensorrt-ensemble/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-tensorrt/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-ts-trace-ensemble/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-ts-trace/config.pbtxt
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/pytorch/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/pytorch/utils.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/docker/build_triton_client.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/download_triton_librispeech.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/execute_all_perf_runs.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/export_model.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/generate_perf_results.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/prepare_model_repository.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/preprocess_triton_librispeech.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_client.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_perf_client.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_server.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/wait_for_triton_server.sh
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/speech_utils.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/tensorrt_io_props_fp16.json
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/tensorrt_io_props_fp32.json
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/triton_librispeech.csv
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/__init__.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/convert_librispeech.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/download_librispeech.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/download_utils.py
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/inference_librispeech.csv
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/librispeech.csv
 create mode 100644 PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/preprocessing_utils.py

diff --git a/PyTorch/contrib/audio/jasper/.keep b/PyTorch/contrib/audio/jasper/.keep
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.dockerignore b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.dockerignore
new file mode 100644
index 0000000000..a620be2e6d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.dockerignore
@@ -0,0 +1,8 @@
+*.pt
+results/
+*__pycache__
+checkpoints/
+.git/
+datasets/
+external/tensorrt-inference-server/
+checkpoints/
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.gitignore b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.gitignore
new file mode 100644
index 0000000000..bb051c475e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.gitignore
@@ -0,0 +1,9 @@
+__pycache__
+*.pt
+results/
+datasets/
+checkpoints/
+
+*.swp
+*.swo
+*.swn
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.keep b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/.keep
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/Dockerfile b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/Dockerfile
new file mode 100644
index 0000000000..8ba48ec348
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/Dockerfile
@@ -0,0 +1,30 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/pytorch:20.10-py3
+FROM ${FROM_IMAGE_NAME}
+
+RUN apt update && apt install -y libsndfile1 && apt install -y sox && rm -rf /var/lib/apt/lists/*
+
+WORKDIR /workspace/jasper
+
+# Install requirements (do this first for better caching)
+COPY requirements.txt .
+RUN conda install -y pyyaml==5.4.1
+RUN pip install --disable-pip-version-check -U -r requirements.txt
+
+RUN pip install --force-reinstall --extra-index-url https://developer.download.nvidia.com/compute/redist nvidia-dali-cuda110==1.2.0
+
+# Copy rest of files
+COPY . .
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/LICENSE b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/LICENSE
new file mode 100644
index 0000000000..2ae5b8195c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/LICENSE
@@ -0,0 +1,203 @@
+   Except where otherwise noted, the following license applies to all files in this repo. 
+        
+                                 Apache License
+                           Version 2.0, January 2004
+                        http://www.apache.org/licenses/
+
+   TERMS AND CONDITIONS FOR USE, REPRODUCTION, AND DISTRIBUTION
+
+   1. Definitions.
+
+      "License" shall mean the terms and conditions for use, reproduction,
+      and distribution as defined by Sections 1 through 9 of this document.
+
+      "Licensor" shall mean the copyright owner or entity authorized by
+      the copyright owner that is granting the License.
+
+      "Legal Entity" shall mean the union of the acting entity and all
+      other entities that control, are controlled by, or are under common
+      control with that entity. For the purposes of this definition,
+      "control" means (i) the power, direct or indirect, to cause the
+      direction or management of such entity, whether by contract or
+      otherwise, or (ii) ownership of fifty percent (50%) or more of the
+      outstanding shares, or (iii) beneficial ownership of such entity.
+
+      "You" (or "Your") shall mean an individual or Legal Entity
+      exercising permissions granted by this License.
+
+      "Source" form shall mean the preferred form for making modifications,
+      including but not limited to software source code, documentation
+      source, and configuration files.
+
+      "Object" form shall mean any form resulting from mechanical
+      transformation or translation of a Source form, including but
+      not limited to compiled object code, generated documentation,
+      and conversions to other media types.
+
+      "Work" shall mean the work of authorship, whether in Source or
+      Object form, made available under the License, as indicated by a
+      copyright notice that is included in or attached to the work
+      (an example is provided in the Appendix below).
+
+      "Derivative Works" shall mean any work, whether in Source or Object
+      form, that is based on (or derived from) the Work and for which the
+      editorial revisions, annotations, elaborations, or other modifications
+      represent, as a whole, an original work of authorship. For the purposes
+      of this License, Derivative Works shall not include works that remain
+      separable from, or merely link (or bind by name) to the interfaces of,
+      the Work and Derivative Works thereof.
+
+      "Contribution" shall mean any work of authorship, including
+      the original version of the Work and any modifications or additions
+      to that Work or Derivative Works thereof, that is intentionally
+      submitted to Licensor for inclusion in the Work by the copyright owner
+      or by an individual or Legal Entity authorized to submit on behalf of
+      the copyright owner. For the purposes of this definition, "submitted"
+      means any form of electronic, verbal, or written communication sent
+      to the Licensor or its representatives, including but not limited to
+      communication on electronic mailing lists, source code control systems,
+      and issue tracking systems that are managed by, or on behalf of, the
+      Licensor for the purpose of discussing and improving the Work, but
+      excluding communication that is conspicuously marked or otherwise
+      designated in writing by the copyright owner as "Not a Contribution."
+
+      "Contributor" shall mean Licensor and any individual or Legal Entity
+      on behalf of whom a Contribution has been received by Licensor and
+      subsequently incorporated within the Work.
+
+   2. Grant of Copyright License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      copyright license to reproduce, prepare Derivative Works of,
+      publicly display, publicly perform, sublicense, and distribute the
+      Work and such Derivative Works in Source or Object form.
+
+   3. Grant of Patent License. Subject to the terms and conditions of
+      this License, each Contributor hereby grants to You a perpetual,
+      worldwide, non-exclusive, no-charge, royalty-free, irrevocable
+      (except as stated in this section) patent license to make, have made,
+      use, offer to sell, sell, import, and otherwise transfer the Work,
+      where such license applies only to those patent claims licensable
+      by such Contributor that are necessarily infringed by their
+      Contribution(s) alone or by combination of their Contribution(s)
+      with the Work to which such Contribution(s) was submitted. If You
+      institute patent litigation against any entity (including a
+      cross-claim or counterclaim in a lawsuit) alleging that the Work
+      or a Contribution incorporated within the Work constitutes direct
+      or contributory patent infringement, then any patent licenses
+      granted to You under this License for that Work shall terminate
+      as of the date such litigation is filed.
+
+   4. Redistribution. You may reproduce and distribute copies of the
+      Work or Derivative Works thereof in any medium, with or without
+      modifications, and in Source or Object form, provided that You
+      meet the following conditions:
+
+      (a) You must give any other recipients of the Work or
+          Derivative Works a copy of this License; and
+
+      (b) You must cause any modified files to carry prominent notices
+          stating that You changed the files; and
+
+      (c) You must retain, in the Source form of any Derivative Works
+          that You distribute, all copyright, patent, trademark, and
+          attribution notices from the Source form of the Work,
+          excluding those notices that do not pertain to any part of
+          the Derivative Works; and
+
+      (d) If the Work includes a "NOTICE" text file as part of its
+          distribution, then any Derivative Works that You distribute must
+          include a readable copy of the attribution notices contained
+          within such NOTICE file, excluding those notices that do not
+          pertain to any part of the Derivative Works, in at least one
+          of the following places: within a NOTICE text file distributed
+          as part of the Derivative Works; within the Source form or
+          documentation, if provided along with the Derivative Works; or,
+          within a display generated by the Derivative Works, if and
+          wherever such third-party notices normally appear. The contents
+          of the NOTICE file are for informational purposes only and
+          do not modify the License. You may add Your own attribution
+          notices within Derivative Works that You distribute, alongside
+          or as an addendum to the NOTICE text from the Work, provided
+          that such additional attribution notices cannot be construed
+          as modifying the License.
+
+      You may add Your own copyright statement to Your modifications and
+      may provide additional or different license terms and conditions
+      for use, reproduction, or distribution of Your modifications, or
+      for any such Derivative Works as a whole, provided Your use,
+      reproduction, and distribution of the Work otherwise complies with
+      the conditions stated in this License.
+
+   5. Submission of Contributions. Unless You explicitly state otherwise,
+      any Contribution intentionally submitted for inclusion in the Work
+      by You to the Licensor shall be under the terms and conditions of
+      this License, without any additional terms or conditions.
+      Notwithstanding the above, nothing herein shall supersede or modify
+      the terms of any separate license agreement you may have executed
+      with Licensor regarding such Contributions.
+
+   6. Trademarks. This License does not grant permission to use the trade
+      names, trademarks, service marks, or product names of the Licensor,
+      except as required for reasonable and customary use in describing the
+      origin of the Work and reproducing the content of the NOTICE file.
+
+   7. Disclaimer of Warranty. Unless required by applicable law or
+      agreed to in writing, Licensor provides the Work (and each
+      Contributor provides its Contributions) on an "AS IS" BASIS,
+      WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or
+      implied, including, without limitation, any warranties or conditions
+      of TITLE, NON-INFRINGEMENT, MERCHANTABILITY, or FITNESS FOR A
+      PARTICULAR PURPOSE. You are solely responsible for determining the
+      appropriateness of using or redistributing the Work and assume any
+      risks associated with Your exercise of permissions under this License.
+
+   8. Limitation of Liability. In no event and under no legal theory,
+      whether in tort (including negligence), contract, or otherwise,
+      unless required by applicable law (such as deliberate and grossly
+      negligent acts) or agreed to in writing, shall any Contributor be
+      liable to You for damages, including any direct, indirect, special,
+      incidental, or consequential damages of any character arising as a
+      result of this License or out of the use or inability to use the
+      Work (including but not limited to damages for loss of goodwill,
+      work stoppage, computer failure or malfunction, or any and all
+      other commercial damages or losses), even if such Contributor
+      has been advised of the possibility of such damages.
+
+   9. Accepting Warranty or Additional Liability. While redistributing
+      the Work or Derivative Works thereof, You may choose to offer,
+      and charge a fee for, acceptance of support, warranty, indemnity,
+      or other liability obligations and/or rights consistent with this
+      License. However, in accepting such obligations, You may act only
+      on Your own behalf and on Your sole responsibility, not on behalf
+      of any other Contributor, and only if You agree to indemnify,
+      defend, and hold each Contributor harmless for any liability
+      incurred by, or claims asserted against, such Contributor by reason
+      of your accepting any such warranty or additional liability.
+
+   END OF TERMS AND CONDITIONS
+
+   APPENDIX: How to apply the Apache License to your work.
+
+      To apply the Apache License to your work, attach the following
+      boilerplate notice, with the fields enclosed by brackets "[]"
+      replaced with your own identifying information. (Don't include
+      the brackets!)  The text should be enclosed in the appropriate
+      comment syntax for the file format. We also recommend that a
+      file or class name and description of purpose be included on the
+      same "printed page" as the copyright notice for easier
+      identification within third-party archives.
+
+   Copyright 2019 NVIDIA Corporation
+
+   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.
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/NOTICE b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/NOTICE
new file mode 100644
index 0000000000..7916839bcc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/NOTICE
@@ -0,0 +1,5 @@
+Jasper in PyTorch
+
+This repository includes source code (in "parts/") from:
+* https://github.com/keithito/tacotron and https://github.com/ryanleary/patter licensed under MIT license.
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/README.md
new file mode 100644
index 0000000000..2bdbc35941
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/README.md
@@ -0,0 +1,49 @@
+# Jasper
+
+This implements training of Jasper on the LibriSpeech dataset.
+
+- Reference implementation:
+```bash
+url=https://github.com/NVIDIA/DeepLearningExamples/tree/master/PyTorch/SpeechRecognition/Jasper
+```
+
+
+## Requirements
+
+- Install PyTorch ([pytorch.org](http://pytorch.org))
+- `pip install -r requirements.txt`
+- Download the LibriSpeech dataset from http://www.openslr.org/12
+   
+
+## Training
+
+- To run the model, you should cd to the directory of test
+
+```bash
+# 1p train full
+bash ./train_full_1p.sh --data_path=xxx
+
+#1p train perf
+bash ./train_performance_1p.sh --data_path=xxx
+
+# 8p train full
+bash ./train_full_8p.sh --data_path=xxx
+
+# 8p train perf
+bash ./train_performance_8p.sh --data_path=xxx
+
+```
+
+
+## Result
+
+| �称   | WER      | 性能       | Epochs |
+| :------: | :------:  | :------: | :------: |  
+| GPU-1p   |    -    |   10       | 1 |     
+| GPU-8p   |  10.73  |   78      | 30 |
+| NPU-1p   |    -    |   4       | 1 |
+| NPU-8p   |  10.89  |   34      | 30 |
+
+
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/README_raw.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/README_raw.md
new file mode 100644
index 0000000000..ad548b22af
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/README_raw.md
@@ -0,0 +1,843 @@
+# Jasper For 
+PyTorch
+
+This repository provides scripts to train the Jasper model to achieve near state of the art accuracy and perform high-performance inference using NVIDIA TensorRT. This repository is tested and maintained by NVIDIA.
+
+## Table Of Contents
+- [Model overview](#model-overview)
+   * [Model architecture](#model-architecture)
+   * [Default configuration](#default-configuration)
+   * [Feature support matrix](#feature-support-matrix)
+       * [Features](#features)
+   * [Mixed precision training](#mixed-precision-training)
+       * [Enabling mixed precision](#enabling-mixed-precision)
+       * [Enabling TF32](#enabling-tf32)
+   * [Glossary](#glossary)
+- [Setup](#setup)
+   * [Requirements](#requirements)
+- [Quick Start Guide](#quick-start-guide)
+- [Advanced](#advanced)
+   * [Scripts and sample code](#scripts-and-sample-code)
+   * [Parameters](#parameters)
+   * [Command-line options](#command-line-options)
+   * [Getting the data](#getting-the-data)
+       * [Dataset guidelines](#dataset-guidelines)
+   * [Training process](#training-process)
+   * [Inference process](#inference-process)
+   * [Evaluation process](#evaluation-process)
+   * [Deploying Jasper using Triton Inference Server](#deploying-jasper-using-triton-inference)
+- [Performance](#performance)
+   * [Benchmarking](#benchmarking)
+       * [Training performance benchmark](#training-performance-benchmark)
+       * [Inference performance benchmark](#inference-performance-benchmark)
+   * [Results](#results)
+       * [Training accuracy results](#training-accuracy-results)
+           * [Training accuracy: NVIDIA DGX A100 (8x A100 80GB)](#training-accuracy-nvidia-dgx-a100-8x-a100-80gb)
+           * [Training accuracy: NVIDIA DGX-1 (8x V100 32GB)](#training-accuracy-nvidia-dgx-1-8x-v100-32gb)
+           * [Training stability test](#training-stability-test)
+       * [Training performance results](#training-performance-results)
+         * [Training performance: NVIDIA DGX A100 (8x A100 80GB)](#training-performance-nvidia-dgx-a100-8x-a100-80gb)
+         * [Training performance: NVIDIA DGX-1 (8x V100 16GB)](#training-performance-nvidia-dgx-1-8x-v100-16gb)
+         * [Training performance: NVIDIA DGX-1 (8x V100 32GB)](#training-performance-nvidia-dgx-1-8x-v100-32gb)
+         * [Training performance: NVIDIA DGX-2 (16x V100 32GB)](#training-performance-nvidia-dgx-2-16x-v100-32gb)
+       * [Inference performance results](#inference-performance-results)
+           * [Inference performance: NVIDIA DGX A100 (1x A100 80GB)](#inference-performance-nvidia-dgx-a100-gpu-1x-a100-80gb)
+           * [Inference performance: NVIDIA DGX-1 (1x V100 16GB)](#inference-performance-nvidia-dgx-1-1x-v100-16gb)
+           * [Inference performance: NVIDIA DGX-1 (1x V100 32GB)](#inference-performance-nvidia-dgx-1-1x-v100-32gb)
+           * [Inference performance: NVIDIA DGX-2 (1x V100 32GB)](#inference-performance-nvidia-dgx-2-1x-v100-32gb)
+           * [Inference performance: NVIDIA T4](#inference-performance-nvidia-t4)
+- [Release notes](#release-notes)
+   * [Changelog](#changelog)
+   * [Known issues](#known-issues)
+
+## Model overview
+This repository provides an implementation of the Jasper model in PyTorch from the paper `Jasper: An End-to-End Convolutional Neural Acoustic Model` [https://arxiv.org/pdf/1904.03288.pdf](https://arxiv.org/pdf/1904.03288.pdf).
+The Jasper model is an end-to-end neural acoustic model for automatic speech recognition (ASR) that provides near state-of-the-art results on LibriSpeech among end-to-end ASR models without any external data. The Jasper architecture of convolutional layers was designed to facilitate fast GPU inference, by allowing whole sub-blocks to be fused into a single GPU kernel. This is important for meeting strict real-time requirements of ASR systems in deployment.
+
+The results of the acoustic model are combined with the results of external language models to get the top-ranked word sequences
+corresponding to a given audio segment. This post-processing step is called decoding.
+
+This repository is a PyTorch implementation of Jasper and provides scripts to train the Jasper 10x5 model with dense residuals from scratch on the [Librispeech](http://www.openslr.org/12) dataset to achieve the greedy decoding results of the original paper.
+The original reference code provides Jasper as part of a research toolkit in TensorFlow [openseq2seq](https://github.com/NVIDIA/OpenSeq2Seq).
+This repository provides a simple implementation of Jasper with scripts for training and replicating the Jasper paper results.
+This includes data preparation scripts, training and inference scripts.
+Both training and inference scripts offer the option to use Automatic Mixed Precision (AMP) to benefit from Tensor Cores for better performance.
+
+In addition to providing the hyperparameters for training a model checkpoint, we publish a thorough inference analysis across different NVIDIA GPU platforms, for example, DGX A100, DGX-1, DGX-2 and T4.
+
+This model is trained with mixed precision using Tensor Cores on Volta, Turing, and the NVIDIA Ampere GPU architectures. Therefore, researchers can get results 3x faster than training without Tensor Cores, while experiencing the benefits of mixed precision training. This model is tested against each NGC monthly container release to ensure consistent accuracy and performance over time.
+
+The original paper takes the output of the Jasper acoustic model and shows results for 3 different decoding variations: greedy decoding, beam search with a 6-gram language model and beam search with further rescoring of the best ranked hypotheses with Transformer XL, which is a neural language model. Beam search and the rescoring with the neural language model scores are run on CPU and result in better word error rates compared to greedy decoding.
+This repository provides instructions to reproduce greedy decoding results. To run beam search or rescoring with TransformerXL, use the following scripts from the [openseq2seq](https://github.com/NVIDIA/OpenSeq2Seq) repository:
+https://github.com/NVIDIA/OpenSeq2Seq/blob/master/scripts/decode.py
+https://github.com/NVIDIA/OpenSeq2Seq/tree/master/external_lm_rescore
+
+### Model architecture
+Details on the model architecture can be found in the paper [Jasper: An End-to-End Convolutional Neural Acoustic Model](https://arxiv.org/pdf/1904.03288.pdf).
+
+|<img src="images/jasper_model.png" width="100%" height="40%"> | <img src="images/jasper_dense_residual.png" width="100%" height="40%">|
+|:---:|:---:|
+|Figure 1: Jasper BxR model: B- number of blocks, R- number of sub-blocks | Figure 2: Jasper Dense Residual |
+
+Jasper is an end-to-end neural acoustic model that is based on convolutions.
+In the audio processing stage, each frame is transformed into mel-scale spectrogram features, which the acoustic model takes as input and outputs a probability distribution over the vocabulary for each frame.
+The acoustic model has a modular block structure and can be parametrized accordingly:
+a Jasper BxR model has B blocks, each consisting of R repeating sub-blocks.
+
+Each sub-block applies the following operations in sequence: 1D-Convolution, Batch Normalization, ReLU activation, and Dropout.
+
+Each block input is connected directly to the last subblock of all following blocks via a residual connection, which is referred to as `dense residual` in the paper.
+Every block differs in kernel size and number of filters, which are increasing in size from the bottom to the top layers.
+Irrespective of the exact block configuration parameters B and R, every Jasper model has four additional convolutional blocks:
+one immediately succeeding the input layer (Prologue) and three at the end of the B blocks (Epilogue).
+
+The Prologue is to decimate the audio signal
+in time in order to process a shorter time sequence for efficiency. The Epilogue with dilation captures a bigger context around an audio time step, which decreases the model word error rate (WER).
+The paper achieves best results with Jasper 10x5 with dense residual connections, which is also the focus of this repository and is in the following referred to as Jasper Large.
+
+### Default configuration
+The following features were implemented in this model:
+
+* GPU-supported feature extraction with data augmentation options [SpecAugment](https://arxiv.org/abs/1904.08779) and [Cutout](https://arxiv.org/pdf/1708.04552.pdf)
+* offline and online [Speed Perturbation](https://www.danielpovey.com/files/2015_interspeech_augmentation.pdf)
+* data-parallel multi-GPU training and evaluation
+* AMP with dynamic loss scaling for Tensor Core training
+* FP16 inference
+
+Competitive training results and analysis is provided for the following Jasper model configuration
+
+|    **Model** | **Number of Blocks** | **Number of Subblocks** | **Max sequence length** | **Number of Parameters** |
+|--------------|----------------------|-------------------------|-------------------------|--------------------------|
+| Jasper Large |                   10 |                       5 |                  16.7 s |                    333 M |
+
+
+### Feature support matrix
+The following features are supported by this model.
+
+| **Feature**   | **Jasper**    |
+|---------------|---------------|
+|[Apex AMP](https://nvidia.github.io/apex/amp.html) | Yes |
+|[Apex DistributedDataParallel](https://nvidia.github.io/apex/parallel.html#apex.parallel.DistributedDataParallel) | Yes |
+
+#### Features
+[Apex AMP](https://nvidia.github.io/apex/amp.html) - a tool that enables Tensor Core-accelerated training. Refer to the [Enabling mixed precision](#enabling-mixed-precision) section for more details.
+
+[Apex
+DistributedDataParallel](https://nvidia.github.io/apex/parallel.html#apex.parallel.DistributedDataParallel) -
+a module wrapper that enables easy multiprocess distributed data parallel
+training, similar to
+[torch.nn.parallel.DistributedDataParallel](https://pytorch.org/docs/stable/nn.html#torch.nn.parallel.DistributedDataParallel).
+`DistributedDataParallel` is optimized for use with
+[NCCL](https://github.com/NVIDIA/nccl). It achieves high performance by
+overlapping communication with computation during `backward()` and bucketing
+smaller gradient transfers to reduce the total number of transfers required.
+
+
+### Mixed precision training
+*Mixed precision* is the combined use of different numerical precisions in a computational method. [Mixed precision](https://arxiv.org/abs/1710.03740) training offers significant computational speedup by performing operations in half-precision format, while storing minimal information in single-precision to retain as much information as possible in critical parts of the network. Since the introduction of [Tensor Cores](https://developer.nvidia.com/tensor-cores) in Volta, and following with both the Turing and Ampere architectures, significant training speedups are experienced by switching to mixed precision -- up to 3x overall speedup on the most arithmetically intense model architectures. Using mixed precision training requires two steps:
+
+1. Porting the model to use the FP16 data type where appropriate.
+2. Adding loss scaling to preserve small gradient values.
+
+The ability to train deep learning networks with lower precision was introduced in the Pascal architecture and first supported in [CUDA 8](https://devblogs.nvidia.com/parallelforall/tag/fp16/) in the NVIDIA Deep Learning SDK.
+
+For information about:
+* How to train using mixed precision, see the[Mixed Precision Training](https://arxiv.org/abs/1710.03740) paper and [Training With Mixed Precision](https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html) documentation.
+* Techniques used for mixed precision training, see the [Mixed-Precision Training of Deep Neural Networks](https://devblogs.nvidia.com/mixed-precision-training-deep-neural-networks/) blog.
+* APEX tools for mixed precision training, see the [NVIDIA Apex: Tools for Easy Mixed-Precision Training in PyTorch](https://devblogs.nvidia.com/apex-pytorch-easy-mixed-precision-training/).
+
+
+#### Enabling mixed precision
+For training, mixed precision can be enabled by setting the flag: `train.py --amp`. When using bash helper scripts:  `scripts/train.sh` `scripts/inference.sh`, etc., mixed precision can be enabled with env variable `AMP=true`.
+
+Mixed precision is enabled in PyTorch by using the Automatic Mixed Precision
+(AMP) library from [APEX](https://github.com/NVIDIA/apex) that casts variables
+to half-precision upon retrieval, while storing variables in single-precision
+format. Furthermore, to preserve small gradient magnitudes in backpropagation,
+a [loss
+scaling](https://docs.nvidia.com/deeplearning/sdk/mixed-precision-training/index.html#lossscaling)
+step must be included when applying gradients. In PyTorch, loss scaling can be
+easily applied by using `scale_loss()` method provided by AMP. The scaling
+value to be used can be
+[dynamic](https://nvidia.github.io/apex/amp.html#apex.amp.initialize) or fixed.
+
+For an in-depth walk through on AMP, check out sample usage
+[here](https://nvidia.github.io/apex/amp.html#). [APEX](https://github.com/NVIDIA/apex) is a PyTorch extension that contains
+utility libraries, such as AMP, which require minimal network code changes to
+leverage Tensor Cores performance.
+
+The following steps were needed to enable mixed precision training in Jasper:
+
+* Import AMP from APEX (file: `train.py`):
+```bash
+from apex import amp
+```
+
+* Initialize AMP and wrap the model and the optimizer
+```bash
+   model, optimizer = amp.initialize(
+     min_loss_scale=1.0,
+     models=model,
+     optimizers=optimizer,
+     opt_level=’O1’)
+
+```
+
+* Apply `scale_loss` context manager
+```bash
+with amp.scale_loss(loss, optimizer) as scaled_loss:
+    scaled_loss.backward()
+```
+
+#### Enabling TF32
+TensorFloat-32 (TF32) is the new math mode in [NVIDIA A100](#https://www.nvidia.com/en-us/data-center/a100/) GPUs for handling the matrix math also called tensor operations. TF32 running on Tensor Cores in A100 GPUs can provide up to 10x speedups compared to single-precision floating-point math (FP32) on Volta GPUs. 
+
+TF32 Tensor Cores can speed up networks using FP32, typically with no loss of accuracy. It is more robust than FP16 for models which require high dynamic range for weights or activations.
+
+For more information, refer to the [TensorFloat-32 in the A100 GPU Accelerates AI Training, HPC up to 20x](#https://blogs.nvidia.com/blog/2020/05/14/tensorfloat-32-precision-format/) blog post.
+
+TF32 is supported in the NVIDIA Ampere GPU architecture and is enabled by default.
+
+
+### Glossary
+**Acoustic model**
+Assigns a probability distribution over a vocabulary of characters given an audio frame.
+
+**Language Model**
+Assigns a probability distribution over a sequence of words. Given a sequence of words, it assigns a probability to the whole sequence.
+
+**Pre-training**
+Training a model on vast amounts of data on the same (or different) task to build general understandings.
+
+**Automatic Speech Recognition (ASR)**
+Uses both acoustic model and language model to output the transcript of an input audio signal.
+
+
+## Setup
+The following section lists the requirements in order to start training and evaluating the Jasper model.
+
+### Requirements
+This repository contains a `Dockerfile` which extends the PyTorch 20.10-py3 NGC container and encapsulates some dependencies. Aside from these dependencies, ensure you have the following components:
+
+* [NVIDIA Docker](https://github.com/NVIDIA/nvidia-docker)
+* [PyTorch 20.10-py3 NGC container](https://ngc.nvidia.com/catalog/containers/nvidia:pytorch)
+-   Supported GPUs:
+    - [NVIDIA Volta architecture](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/)
+    - [NVIDIA Turing architecture](https://www.nvidia.com/en-us/geforce/turing/)
+    - [NVIDIA Ampere architecture](https://www.nvidia.com/en-us/data-center/nvidia-ampere-gpu-architecture/)
+
+Further required python packages are listed in `requirements.txt`, which are automatically installed with the Docker container built. To manually install them, run
+```bash
+pip install -r requirements.txt
+```
+
+For more information about how to get started with NGC containers, see the following sections from the NVIDIA GPU Cloud Documentation and the Deep Learning Documentation:
+
+* [Getting Started Using NVIDIA GPU Cloud](https://docs.nvidia.com/ngc/ngc-getting-started-guide/index.html)
+* [Accessing And Pulling From The NGC Container Registry](https://docs.nvidia.com/deeplearning/dgx/user-guide/index.html#accessing_registry)
+* [Running PyTorch](https://docs.nvidia.com/deeplearning/dgx/pytorch-release-notes/running.html#running)
+
+For those unable to use the PyTorch NGC container, to set up the required environment or create your own container, see the versioned [NVIDIA Container Support Matrix](https://docs.nvidia.com/deeplearning/dgx/support-matrix/index.html).
+
+
+## Quick Start Guide
+
+To train your model using mixed or TF32 precision with Tensor Cores or using FP32, perform the following steps using the default parameters of the Jasper model on the Librispeech dataset. For details concerning training and inference, see [Advanced](#Advanced) section.
+
+1. Clone the repository.
+```bash
+git clone https://github.com/NVIDIA/DeepLearningExamples
+cd DeepLearningExamples/PyTorch/SpeechRecognition/Jasper
+```
+2. Build the Jasper PyTorch container.
+
+Running the following scripts will build and launch the container which contains all the required dependencies for data download and processing as well as training and inference of the model.
+
+```bash
+bash scripts/docker/build.sh
+```
+
+3. Start an interactive session in the NGC container to run data download/training/inference
+
+```bash
+bash scripts/docker/launch.sh <DATA_DIR> <CHECKPOINT_DIR> <OUTPUT_DIR>
+```
+Within the container, the contents of this repository will be copied to the `/workspace/jasper` directory. The `/datasets`, `/checkpoints`, `/results` directories are mounted as volumes
+and mapped to the corresponding directories `<DATA_DIR>`, `<CHECKPOINT_DIR>`, `<OUTPUT_DIR>` on the host.
+
+4. Download and preprocess the dataset.
+
+No GPU is required for data download and preprocessing. Therefore, if GPU usage is a limited resource, launch the container for this section on a CPU machine by following Steps 2 and 3.
+
+Note: Downloading and preprocessing the dataset requires 500GB of free disk space and can take several hours to complete.
+
+This repository provides scripts to download, and extract the following datasets:
+
+* LibriSpeech [http://www.openslr.org/12](http://www.openslr.org/12)
+
+LibriSpeech contains 1000 hours of 16kHz read English speech derived from public domain audiobooks from LibriVox project and has been carefully segmented and aligned. For more information, see the [LIBRISPEECH: AN ASR CORPUS BASED ON PUBLIC DOMAIN AUDIO BOOKS](http://www.danielpovey.com/files/2015_icassp_librispeech.pdf) paper.
+
+Inside the container, download and extract the datasets into the required format for later training and inference:
+```bash
+bash scripts/download_librispeech.sh
+```
+Once the data download is complete, the following folders should exist:
+```bash
+datasets/LibriSpeech/
+├── dev-clean
+├── dev-other
+├── test-clean
+├── test-other
+├── train-clean-100
+├── train-clean-360
+└── train-other-500
+```
+
+Since `/datasets/` is mounted to `<DATA_DIR>` on the host (see Step 3),  once the dataset is downloaded it will be accessible from outside of the container at `<DATA_DIR>/LibriSpeech`.
+
+
+Next, convert the data into WAV files:
+```bash
+bash scripts/preprocess_librispeech.sh
+```
+Once the data is converted, the following additional files and folders should exist:
+```bash
+datasets/LibriSpeech/
+├── dev-clean-wav
+├── dev-other-wav
+├── librispeech-train-clean-100-wav.json
+├── librispeech-train-clean-360-wav.json
+├── librispeech-train-other-500-wav.json
+├── librispeech-dev-clean-wav.json
+├── librispeech-dev-other-wav.json
+├── librispeech-test-clean-wav.json
+├── librispeech-test-other-wav.json
+├── test-clean-wav
+├── test-other-wav
+├── train-clean-100-wav
+├── train-clean-360-wav
+└── train-other-500-wav
+```
+
+The DALI data pre-processing pipeline, which is enabled by default, performs speed perturbation on-line during training.
+Without DALI, on-line speed perturbation might slow down the training.
+If you wish to disable DALI, speed perturbation can be computed off-line with:
+```bash
+SPEEDS="0.9 1.1" bash scripts/preprocess_librispeech.sh
+```
+
+5. Start training.
+
+Inside the container, use the following script to start training.
+Make sure the downloaded and preprocessed dataset is located at `<DATA_DIR>/LibriSpeech` on the host (see Step 3), which corresponds to `/datasets/LibriSpeech` inside the container.
+
+```bash
+[OPTION1=value1 OPTION2=value2 ...] bash scripts/train.sh
+```
+By default automatic precision is disabled, batch size is 64 over two gradient accumulation steps, and the recipe is run on a total of 8 GPUs. The hyperparameters are tuned for a GPU with at least 32GB of memory and will require adjustment for 16GB GPUs (e.g., by lowering batch size and using more gradient accumulation steps).
+
+Options are being passed as environment variables. More details on available options can be found in [Parameters](#parameters) and [Training process](#training-process).
+
+6. Start validation/evaluation.
+
+Inside the container, use the following script to run evaluation.
+ Make sure the downloaded and preprocessed dataset is located at `<DATA_DIR>/LibriSpeech` on the host (see Step 3), which corresponds to `/datasets/LibriSpeech` inside the container.
+```bash
+[OPTION1=value1 OPTION2=value2 ...] bash scripts/evaluation.sh [OPTIONS]
+```
+By default, this will use full precision, a batch size of 64 and run on a single GPU.
+
+Options are being passed as environment variables. More details on available options can be found in [Parameters](#parameters) and [Evaluation process](#evaluation-process).
+
+
+7. Start inference/predictions.
+
+Inside the container, use the following script to run inference.
+ Make sure the downloaded and preprocessed dataset is located at `<DATA_DIR>/LibriSpeech` on the host (see Step 3), which corresponds to `/datasets/LibriSpeech` inside the container.
+A pretrained model checkpoint can be downloaded from [NGC model repository](https://ngc.nvidia.com/catalog/models/nvidia:jasperpyt_fp16).
+
+```bash
+[OPTION1=value1 OPTION2=value2 ...] bash scripts/inference.sh
+```
+By default this will use single precision, a batch size of 64 and run on a single GPU.
+
+Options are being passed as environment variables. More details on available options can be found in [Parameters](#parameters) and [Inference process](#inference-process).
+
+
+## Advanced
+
+The following sections provide greater details of the dataset, running training and inference, and getting training and inference results.
+
+
+### Scripts and sample code
+In the `root` directory, the most important files are:
+```
+jasper
+├── common        # data pre-processing, logging, etc.
+├── configs       # model configurations
+├── Dockerfile    # container with the basic set of dependencies to run Jasper
+├── inference.py  # entry point for inference
+├── jasper        # model-specific code
+├── notebooks     # jupyter notebooks and example audio files
+├── scripts       # one-click scripts required for running various supported functionalities
+│   ├── docker                     # contains the scripts for building and launching the container
+│   ├── download_librispeech.sh    # downloads LibriSpeech dataset
+│   ├── evaluation.sh              # runs evaluation using the `inference.py` script
+│   ├── inference_benchmark.sh     # runs the inference benchmark using the `inference_benchmark.py` script
+│   ├── inference.sh               # runs inference using the `inference.py` script
+│   ├── preprocess_librispeech.sh  # preprocess LibriSpeech raw data files for training and inference
+│   ├── train_benchmark.sh         # runs the training performance benchmark using the `train.py` script
+│   └── train.sh                   # runs training using the `train.py` script
+├── train.py      # entry point for training
+├── triton        # example of inference using Triton Inference Server
+└── utils         # data downloading and common routines
+```
+
+### Parameters
+
+Parameters could be set as env variables, or passed as positional arguments.
+
+The complete list of available parameters for `scripts/train.sh` script contains:
+```bash
+DATA_DIR:                directory of dataset. (default: '/datasets/LibriSpeech')
+MODEL_CONFIG:            relative path to model configuration. (default: 'configs/jasper10x5dr_speedp-online_speca.yaml')
+OUTPUT_DIR:              directory for results, logs, and created checkpoints. (default: '/results')
+CHECKPOINT:              a specific model checkpoint to continue training from. To resume training from the last checkpoint, see the RESUME option.
+RESUME:                  resume training from the last checkpoint found in OUTPUT_DIR, or from scratch if there are no checkpoints (default: true)
+CUDNN_BENCHMARK:         boolean that indicates whether to enable cudnn benchmark mode for using more optimized kernels. (default: true)
+NUM_GPUS:                number of GPUs to use. (default: 8)
+AMP:                     if set to `true`, enables automatic mixed precision (default: false)
+BATCH_SIZE:              effective data batch size. The real batch size per GPU might be lower, if gradient accumulation is enabled (default: 64)
+GRAD_ACCUMULATION_STEPS: number of gradient accumulation steps until optimizer updates weights. (default: 2)
+LEARNING_RATE:           initial learning rate. (default: 0.01)
+MIN_LEARNING_RATE:       minimum learning rate, despite LR scheduling (default: 1e-5)
+LR_POLICY:               how to decay LR (default: exponential)
+LR_EXP_GAMMA:            decay factor for the exponential LR schedule (default: 0.981)
+EMA:                     decay factor for exponential averages of checkpoints (default: 0.999)
+SEED:                    seed for random number generator and used for ensuring reproducibility. (default: 0)
+EPOCHS:                  number of training epochs. (default: 440)
+WARMUP_EPOCHS:           number of initial epoch of linearly increasing LR. (default: 2)
+HOLD_EPOCHS:             number of epochs to hold maximum LR after warmup. (default: 140)
+SAVE_FREQUENCY:          number of epochs between saving the model to disk. (default: 10)
+EPOCHS_THIS_JOB:         run training for this number of epochs. Does not affect LR schedule like the EPOCHS parameter. (default: 0)
+DALI_DEVICE:             device to run the DALI pipeline on for calculation of filterbanks. Valid choices: cpu, gpu, none. (default: gpu)
+PAD_TO_MAX_DURATION:     pad all sequences with zeros to maximum length. (default: false)
+EVAL_FREQUENCY:          number of steps between evaluations on the validation set. (default: 544)
+PREDICTION_FREQUENCY:    the number of steps between writing a sample prediction to stdout. (default: 544)
+TRAIN_MANIFESTS:         lists of .json training set files
+VAL_MANIFESTS:           lists of .json validation set files
+
+```
+
+The complete list of available parameters for `scripts/inference.sh` script contains:
+```bash
+DATA_DIR:            directory of dataset. (default: '/datasets/LibriSpeech')
+MODEL_CONFIG:        model configuration. (default: 'configs/jasper10x5dr_speedp-online_speca.yaml')
+OUTPUT_DIR:          directory for results and logs. (default: '/results')
+CHECKPOINT:          model checkpoint path. (required)
+DATASET:             name of the LibriSpeech subset to use. (default: 'dev-clean')
+LOG_FILE:            path to the DLLogger .json logfile. (default: '')
+CUDNN_BENCHMARK:     enable cudnn benchmark mode for using more optimized kernels. (default: false)
+MAX_DURATION:        filter out recordings shorter then MAX_DURATION seconds. (default: "")
+PAD_TO_MAX_DURATION: pad all sequences with zeros to maximum length. (default: false)
+PAD_LEADING:         pad every batch with leading zeros to counteract conv shifts of the field of view. (default: 16)
+NUM_GPUS:            number of GPUs to use. Note that with > 1 GPUs WER results might be inaccurate due to the batching policy. (default: 1)
+NUM_STEPS:           number of batches to evaluate, loop the dataset if necessary. (default: 0)
+NUM_WARMUP_STEPS:    number of initial steps before measuring performance. (default: 0)
+AMP:                 enable FP16 inference with AMP. (default: false)
+BATCH_SIZE:          data batch size. (default: 64)
+EMA:                 Attempt to load exponentially averaged weights from a checkpoint. (default: true)
+SEED:                seed for random number generator and used for ensuring reproducibility. (default: 0)
+DALI_DEVICE:         device to run the DALI pipeline on for calculation of filterbanks. Valid choices: cpu, gpu, none. (default: gpu)
+CPU:                 run inference on CPU. (default: false)
+LOGITS_FILE:         dump logit matrices to a file. (default: "")
+PREDICTION_FILE:     save predictions to a file. (default: "${OUTPUT_DIR}/${DATASET}.predictions")
+```
+
+The complete list of available parameters for `scripts/evaluation.sh` is the same as for `scripts/inference.sh` except for the few default changes.
+```bash
+PREDICTION_FILE: (default: "")
+DATASET:         (default: "test-other")
+```
+
+The `scripts/inference_benchmark.sh` script pads all input to a fixed duration and computes the mean, 90%, 95%, 99% percentile of latency for the specified number of inference steps. Latency is measured in milliseconds per batch. The `scripts/inference_benchmark.sh` measures latency for a single GPU and loops over a number of batch sizes and durations. It extends  `scripts/inference.sh`, and changes the defaults with:
+```bash
+BATCH_SIZE_SEQ:      batch sizes to measure on. (default: "1 2 4 8 16")
+MAX_DURATION_SEQ:    input durations (in seconds) to measure on (default: "2 7 16.7")
+CUDNN_BENCHMARK:     (default: true)
+PAD_TO_MAX_DURATION: (default: true)
+PAD_LEADING:         (default: 0)
+NUM_WARMUP_STEPS:    (default: 10)
+NUM_STEPS:           (default: 500)
+DALI_DEVICE:         (default: cpu)
+```
+
+The `scripts/train_benchmark.sh` script pads all input to the same length according to the input argument `MAX_DURATION` and measures average training latency and throughput performance. Latency is measured in seconds per batch, throughput in sequences per second.
+Training performance is measured with on-line speed perturbation and cuDNN benchmark mode enabled.
+The script `scripts/train_benchmark.sh` loops over a number of batch sizes and GPU counts.
+It extends `scripts/train.sh`, and the complete list of available parameters for `scripts/train_benchmark.sh` script contains:
+```bash
+BATCH_SIZE_SEQ:          batch sizes to measure on. (default: "1 2 4 8 16")
+NUM_GPUS_SEQ:            number of GPUs to run the training on. (default: "1 4 8")
+MODEL_CONFIG:            (default: "configs/jasper10x5dr_speedp-online_train-benchmark.yaml")
+TRAIN_MANIFESTS:         (default: "$DATA_DIR/librispeech-train-clean-100-wav.json")
+RESUME:                  (default: false)
+EPOCHS_THIS_JOB:         (default: 2)
+EPOCHS:                  (default: 100000)
+SAVE_FREQUENCY:          (default: 100000)
+EVAL_FREQUENCY:          (default: 100000)
+GRAD_ACCUMULATION_STEPS: (default: 1)
+PAD_TO_MAX_DURATION:     (default: true)
+EMA:                     (default: 0)
+```
+
+### Command-line options
+To see the full list of available options and their descriptions, use the `-h` or `--help` command-line option with the Python file, for example:
+
+```bash
+python train.py --help
+python inference.py --help
+```
+
+### Getting the data
+The Jasper model was trained on the LibriSpeech dataset. We use the concatenation of `train-clean-100`, `train-clean-360` and `train-other-500` for training and `dev-clean` for validation.
+
+This repository contains the `scripts/download_librispeech.sh` and `scripts/preprocess_librispeech.sh` scripts which will automatically download and preprocess the training, test and development datasets. By default, data will be downloaded to the `/datasets/LibriSpeech` directory, a minimum of 250GB free space is required for download and preprocessing, the final preprocessed dataset is approximately 100GB. With offline speed perturbation, the dataset will be about 3x larger.
+
+#### Dataset guidelines
+The `scripts/preprocess_librispeech.sh` script converts the input audio files to WAV format with a sample rate of 16kHz, target transcripts are stripped from whitespace characters, then lower-cased. For `train-clean-100`, `train-clean-360` and `train-other-500`. It can optionally create speed perturbed versions with rates of 0.9 and 1.1 for data augmentation. In the current version, those augmentations are applied on-line with the DALI pipeline without any impact on training time.
+
+After preprocessing, the script creates JSON files with output file paths, sample rate, target transcript and other metadata. These JSON files are used by the training script to identify training and validation datasets.
+
+The Jasper model was tuned on audio signals with a sample rate of 16kHz, if you wish to use a different sampling rate then some hyperparameters might need to be changed - specifically window size and step size.
+
+
+### Training process
+
+The training is performed using `train.py` script along with parameters defined in  `scripts/train.sh`
+The `scripts/train.sh` script runs a job on a single node that trains the Jasper model from scratch using LibriSpeech as training data. To make training more efficient, we discard audio samples longer than 16.7 seconds from the training dataset, the total number of these samples is less than 1%. Such filtering does not degrade accuracy, but it allows us to decrease the number of time steps in a batch, which requires less GPU memory and increases training speed.
+Apart from the default arguments as listed in the [Parameters](#parameters) section, by default the training script:
+
+* Runs on 8 GPUs with at least 32GB of memory and training/evaluation batch size 64, split over two gradient accumulation steps
+* Uses TF32 precision (A100 GPU) or FP32 (other GPUs)
+* Trains on the concatenation of all 3 LibriSpeech training datasets and evaluates on the LibriSpeech dev-clean dataset
+* Maintains an exponential moving average of parameters for evaluation
+* Has cudnn benchmark enabled
+* Runs for 440 epochs
+* Uses an initial learning rate of 0.01 and an exponential learning rate decay
+* Saves a checkpoint every 10 epochs
+* Automatically removes old checkpoints and preserves milestone checkpoints
+* Runs evaluation on the development dataset every 544 iterations and at the end of training
+* Maintains a separate checkpoint with the lowest WER on development set
+* Prints out training progress every iteration to stdout
+* Creates a DLLogger logfile and a Tensorboard log
+* Calculates speed perturbation on-line during training
+* Uses SpecAugment in data pre-processing
+* Filters out audio samples longer than 16.7 seconds
+* Pads each batch so its length would be divisible by 16
+* Uses masked convolutions and dense residuals as described in the paper
+* Uses weight decay of 0.001
+* Uses [Novograd](https://arxiv.org/pdf/1905.11286.pdf) as optimizer with betas=(0.95, 0)
+
+Enabling AMP permits batch size 64 with one gradient accumulation step. In the current setup it will improve upon the greedy WER [Results](#results) of the Jasper paper on a DGX-1 with 32GB V100 GPUs.
+
+### Inference process
+Inference is performed using the `inference.py` script along with parameters defined in `scripts/inference.sh`.
+The `scripts/inference.sh` script runs the job on a single GPU, taking a pre-trained Jasper model checkpoint and running it on the specified dataset.
+Apart from the default arguments as listed in the [Parameters](#parameters) section by default the inference script:
+
+* Evaluates on the LibriSpeech dev-clean dataset
+* Uses a batch size of 64
+* Runs for 1 epoch and prints out the final word error rate
+* Creates a log file with progress and results which will be stored in the results folder
+* Pads each batch so its length would be divisible by 16
+* Does not use data augmentation
+* Does greedy decoding and saves the transcription in the results folder
+* Has the option to save the model output tensors for more complex decoding, for example, beam search
+* Has cudnn benchmark disabled
+
+### Evaluation process
+Evaluation is performed using the `inference.py` script along with parameters defined in `scripts/evaluation.sh`.
+The setup is similar to `scripts/inference.sh`, with two differences:
+
+* Evaluates the LibriSpeech test-other dataset
+* Model outputs are not saved
+
+### Deploying Jasper using Triton Inference Server
+The NVIDIA Triton Inference Server provides a datacenter and cloud inferencing solution optimized for NVIDIA GPUs. The server provides an inference service via an HTTP or gRPC endpoint, allowing remote clients to request inferencing for any number of GPU or CPU models being managed by the server.
+More information on how to perform inference using Triton Inference Server with different model backends can be found in the subfolder [./triton/README.md](triton/README.md)
+
+
+## Performance
+
+The performance measurements in this document were conducted at the time of publication and may not reflect the performance achieved from NVIDIA’s latest software release. For the most up-to-date performance measurements, go to [NVIDIA Data Center Deep Learning Product Performance](https://developer.nvidia.com/deep-learning-performance-training-inference).
+
+### Benchmarking
+The following section shows how to run benchmarks measuring the model performance in training and inference modes.
+
+#### Training performance benchmark
+To benchmark the training performance in a specific setting on the `train-clean-100` subset of LibriSpeech, run:
+
+```bash
+BATCH_SIZE_SEQ=<BATCH_SIZES> NUM_GPUS_SEQ=<NUMS_OF_GPUS> bash scripts/train_benchmark.sh
+```
+
+By default, this script runs 2 epochs on the configuration `configs/jasper10x5dr_speedp-online_train-benchmark.yaml`,
+which applies gentle speed perturbation that does not change the length of the output, enabling immediate stabilization of training step times in the cuDNN benchmark mode. The script runs benchmarks on batch sizes 32 on 1, 4, and 8 GPUs, and requires a 8x 32GB GPU machine.
+
+#### Inference performance benchmark
+To benchmark the inference performance on a specific batch size and audio length, run:
+
+```bash
+BATCH_SIZE_SEQ=<BATCH_SIZES> MAX_DURATION_SEQ=<DURATIONS> bash scripts/inference_benchmark.sh
+```
+By default, the script runs on a single GPU and evaluates on the dataset limited to utterances shorter than MAX_DURATION. It uses the model configuration `configs/jasper10x5dr_speedp-online_speca.yaml`.
+
+
+### Results
+The following sections provide details on how we achieved our performance and accuracy in training and inference.
+All results are trained on 960 hours of LibriSpeech with a maximum audio length of 16.7s. The training is evaluated
+on LibriSpeech dev-clean, dev-other, test-clean, test-other. Checkpoints for evaluation are being chosen based on their
+word error rate on dev-clean.
+
+#### Training accuracy results
+
+##### Training accuracy: NVIDIA DGX A100 (8x A100 80GB)
+Our results were obtained by running the `scripts/train.sh` training script in the PyTorch 20.10-py3 NGC container with NVIDIA DGX A100 with (8x A100 80GB) GPUs.
+The following table reports the word error rate (WER) of the acoustic model with greedy decoding on all LibriSpeech dev and test datasets for mixed precision training.
+
+| Number of GPUs | Batch size per GPU | Precision | dev-clean WER | dev-other WER | test-clean WER | test-other WER | Time to train |
+|-----|-----|-------|-------|-------|------|-------|------|
+|   8 |  64 | mixed |  3.20 |  9.78 | 3.41 |  9.71 | 70 h |
+
+##### Training accuracy: NVIDIA DGX-1 (8x V100 32GB)
+Our results were obtained by running the `scripts/train.sh` training script in the PyTorch 20.10-py3 NGC container with NVIDIA DGX-1 with (8x V100 32GB) GPUs.
+The following table reports the word error rate (WER) of the acoustic model with greedy decoding on all LibriSpeech dev and test datasets for mixed precision training.
+
+| Number of GPUs | Batch size per GPU | Precision | dev-clean WER | dev-other WER | test-clean WER | test-other WER | Time to train |
+|-----|-----|-------|-------|-------|------|-------|-------|
+|   8 |  64 | mixed |  3.26 | 10.00 | 3.54 |  9.80 | 130 h |
+
+We show the best of 5 runs (mixed precision) and 2 runs (FP32) chosen based on dev-clean WER. For FP32, two gradient accumulation steps have been used.
+
+##### Training stability test
+The following table compares greedy decoding word error rates across 8 different training runs with different seeds for mixed precision training.
+
+| DGX A100 80GB, FP16, 8x GPU |   Seed #1 |   Seed #2 |   Seed #3 |   Seed #4 |   Seed #5 |   Seed #6 |   Seed #7 |   Seed #8 |   Mean |   Std |
+|-----------:|----------:|----------:|----------:|----------:|----------:|----------:|----------:|----------:|-------:|------:|
+| dev-clean  |      3.46 |      3.55 |      3.45 |      3.44 |      3.25 |      3.34 |      3.20 |      3.40 |   3.39 |  0.11 |
+| dev-other  |     10.30 |     10.77 |     10.36 |     10.26 |      9.99 |     10.18 |      9.78 |     10.32 |  10.25 |  0.27 |
+| test-clean |      3.84 |      3.81 |      3.66 |      3.64 |      3.58 |      3.55 |      3.41 |      3.73 |   3.65 |  0.13 |
+| test-other |     10.61 |     10.52 |     10.49 |     10.47 |      9.89 |     10.09 |      9.71 |     10.26 |  10.26 |  0.31 |
+
+
+| DGX-1 32GB, FP16, 8x GPU |   Seed #1 |   Seed #2 |   Seed #3 |   Seed #4 |   Seed #5 |   Seed #6 |   Seed #7 |   Seed #8 |   Mean |   Std |
+|-----------:|----------:|----------:|----------:|----------:|----------:|----------:|----------:|----------:|-------:|------:|
+| dev-clean  |      3.31 |      3.31 |      3.26 |      3.44 |      3.40 |      3.35 |      3.36 |      3.28 |   3.34 |  0.06 |
+| dev-other  |     10.02 |     10.01 |     10.00 |     10.06 |     10.05 |     10.03 |     10.10 |     10.04 |  10.04 |  0.03 |
+| test-clean |      3.49 |      3.50 |      3.54 |      3.61 |      3.57 |      3.58 |      3.48 |      3.51 |   3.54 |  0.04 |
+| test-other |     10.11 |     10.14 |      9.80 |     10.09 |     10.17 |      9.99 |      9.86 |     10.00 |  10.02 |  0.13 |
+
+#### Training performance results
+Our results were obtained by running the `scripts/train.sh` training script in the PyTorch 20.10-py3 NGC container. Performance (in sequences per second) is the steady-state throughput.
+
+##### Training performance: NVIDIA DGX A100 (8x A100 80GB)
+| Batch size / GPU | GPUs | Throughput - TF32 | Throughput - mixed precision | Throughput speedup (TF32 to mixed precision) | Weak scaling - TF32 | Weak scaling - mixed precision |
+|----:|----:|-------:|-------:|-----:|-----:|-----:|
+|  32 |   1 |  42.18 |  64.32 | 1.52 | 1.00 | 1.00 |
+|  32 |   4 | 157.49 | 239.23 | 1.52 | 3.73 | 3.72 |
+|  32 |   8 | 310.10 | 470.09 | 1.52 | 7.35 | 7.31 |
+|  64 |   1 |  49.64 |  75.59 | 1.52 | 1.00 | 1.00 |
+|  64 |   4 | 192.66 | 289.16 | 1.50 | 3.88 | 3.83 |
+|  64 |   8 | 371.41 | 547.91 | 1.48 | 7.48 | 7.25 |
+
+Note: Mixed precision permits higher batch sizes during training. We report the maximum batch sizes (as powers of 2), which are allowed without gradient accumulation.
+
+To achieve these same results, follow the [Quick Start Guide](#quick-start-guide) outlined above.
+
+##### Training performance: NVIDIA DGX-1 (8x V100 16GB)
+| Batch size / GPU | GPUs | Throughput - FP32 | Throughput - mixed precision | Throughput speedup (FP32 to mixed precision) | Weak scaling - FP32 | Weak scaling - mixed precision |
+|----:|----:|------:|-------:|-----:|-----:|-----:|
+|  16 |   1 | 10.71 |  27.87 | 2.60 | 1.00 | 1.00 |
+|  16 |   4 | 40.28 |  99.80 | 2.48 | 3.76 | 3.58 |
+|  16 |   8 | 78.23 | 193.89 | 2.48 | 7.30 | 6.96 |
+
+Note: Mixed precision permits higher batch sizes during training. We report the maximum batch sizes (as powers of 2), which are allowed without gradient accumulation.
+
+To achieve these same results, follow the [Quick Start Guide](#quick-start-guide) outlined above.
+
+##### Training performance: NVIDIA DGX-1 (8x V100 32GB)
+| Batch size / GPU | GPUs | Throughput - FP32 | Throughput - mixed precision | Throughput speedup (FP32 to mixed precision) | Weak scaling - FP32 | Weak scaling - mixed precision |
+|----:|----:|------:|-------:|-----:|-----:|-----:|
+|  32 |   1 | 12.22 |  34.08 | 2.79 | 1.00 | 1.00 |
+|  32 |   4 | 46.97 | 128.39 | 2.73 | 3.84 | 3.77 |
+|  32 |   8 | 92.44 | 249.00 | 2.69 | 7.57 | 7.31 |
+|  64 |   1 |   N/A |  39.30 |  N/A |  N/A | 1.00 |
+|  64 |   4 |   N/A | 150.18 |  N/A |  N/A | 3.82 |
+|  64 |   8 |   N/A | 282.68 |  N/A |  N/A | 7.19 |
+
+Note: Mixed precision permits higher batch sizes during training. We report the maximum batch sizes (as powers of 2), which are allowed without gradient accumulation.
+
+To achieve these same results, follow the [Quick Start Guide](#quick-start-guide) outlined above.
+
+##### Training performance: NVIDIA DGX-2 (16x V100 32GB)
+| Batch size / GPU | GPUs | Throughput - FP32 | Throughput - mixed precision | Throughput speedup (FP32 to mixed precision) | Weak scaling - FP32 | Weak scaling - mixed precision |
+|----:|----:|-------:|-------:|-----:|------:|------:|
+|  32 |   1 |  13.46 |  38.94 | 2.89 |  1.00 |  1.00 |
+|  32 |   4 |  51.38 | 143.44 | 2.79 |  3.82 |  3.68 |
+|  32 |   8 | 100.54 | 280.48 | 2.79 |  7.47 |  7.20 |
+|  32 |  16 | 188.14 | 515.90 | 2.74 | 13.98 | 13.25 |
+|  64 |   1 |    N/A |  43.86 |  N/A |   N/A |  1.00 |
+|  64 |   4 |    N/A | 165.27 |  N/A |   N/A |  3.77 |
+|  64 |   8 |    N/A | 318.10 |  N/A |   N/A |  7.25 |
+|  64 |  16 |    N/A | 567.47 |  N/A |   N/A | 12.94 |
+
+Note: Mixed precision permits higher batch sizes during training. We report the maximum batch sizes (as powers of 2), which are allowed without gradient accumulation.
+
+To achieve these same results, follow the [Quick Start Guide](#quick-start-guide) outlined above.
+
+
+#### Inference performance results
+Our results were obtained by running the `scripts/inference_benchmark.sh` script in the PyTorch 20.10-py3 NGC container on NVIDIA DGX A100, DGX-1, DGX-2 and T4 on a single GPU. Performance numbers (latency in milliseconds per batch) were averaged over 500 iterations.
+
+##### Inference performance: NVIDIA DGX A100 (1x A100 80GB)
+|  |  | FP16 Latency (ms) Percentiles |  |  |  | TF32 Latency (ms) Percentiles |  |  |  | FP16/TF32 speed up |
+|-----:|---------------:|------:|------:|------:|------:|------:|------:|-------:|------:|------:|
+|   BS |   Duration (s) |   90% |   95% |   99% |   Avg |   90% |   95% |    99% |   Avg |   Avg |
+|    1 |            2.0 | 32.40 | 32.50 | 32.82 | 32.30 | 33.30 | 33.64 |  34.65 | 33.25 |  1.03 |
+|    2 |            2.0 | 32.90 | 33.51 | 34.35 | 32.69 | 34.48 | 34.65 |  35.66 | 34.27 |  1.05 |
+|    4 |            2.0 | 32.85 | 33.01 | 33.89 | 32.60 | 34.09 | 34.46 |  35.22 | 34.00 |  1.04 |
+|    8 |            2.0 | 35.51 | 35.89 | 37.10 | 35.33 | 34.86 | 35.36 |  36.08 | 34.45 |  0.98 |
+|   16 |            2.0 | 36.00 | 36.57 | 37.40 | 35.77 | 43.83 | 44.12 |  44.77 | 43.39 |  1.21 |
+|    1 |            7.0 | 33.50 | 33.99 | 34.91 | 33.03 | 33.83 | 34.25 |  34.95 | 33.70 |  1.02 |
+|    2 |            7.0 | 34.43 | 34.89 | 35.72 | 34.22 | 34.41 | 34.73 |  35.69 | 34.28 |  1.00 |
+|    4 |            7.0 | 34.30 | 34.59 | 35.43 | 34.07 | 37.95 | 38.18 |  38.87 | 37.55 |  1.10 |
+|    8 |            7.0 | 35.98 | 36.28 | 37.11 | 35.28 | 44.64 | 44.79 |  45.37 | 44.29 |  1.26 |
+|   16 |            7.0 | 39.86 | 40.08 | 41.16 | 39.33 | 55.17 | 55.46 |  57.24 | 54.56 |  1.39 |
+|    1 |           16.7 | 35.20 | 35.80 | 38.71 | 34.36 | 35.36 | 35.76 |  36.55 | 34.64 |  1.01 |
+|    2 |           16.7 | 35.40 | 35.81 | 36.50 | 34.76 | 36.34 | 36.53 |  37.40 | 35.87 |  1.03 |
+|    4 |           16.7 | 36.01 | 36.38 | 37.37 | 35.57 | 44.69 | 45.09 |  45.88 | 43.92 |  1.23 |
+|    8 |           16.7 | 41.48 | 41.78 | 44.22 | 40.69 | 58.57 | 58.74 |  59.62 | 58.11 |  1.43 |
+|   16 |           16.7 | 61.37 | 61.93 | 66.32 | 60.92 | 97.33 | 97.71 | 100.04 | 96.56 |  1.59 |
+
+To achieve these same results, follow the [Quick Start Guide](#quick-start-guide) outlined above.
+
+##### Inference performance: NVIDIA DGX-1 (1x V100 16GB)
+|  |  | FP16 Latency (ms) Percentiles |  |  |  | FP32 Latency (ms) Percentiles |  |  |  | FP16/FP32 speed up |
+|-----:|---------------:|-------:|-------:|-------:|-------:|-------:|-------:|-------:|-------:|------:|
+|   BS |   Duration (s) |    90% |    95% |    99% |    Avg |    90% |    95% |    99% |    Avg |   Avg |
+|    1 |            2.0 |  45.42 |  45.62 |  49.54 |  45.02 |  48.83 |  48.99 |  51.66 |  48.44 |  1.08 |
+|    2 |            2.0 |  50.31 |  50.53 |  53.66 |  49.10 |  49.87 |  50.04 |  52.99 |  49.41 |  1.01 |
+|    4 |            2.0 |  49.17 |  49.48 |  52.13 |  48.73 |  52.92 |  53.21 |  55.28 |  52.31 |  1.07 |
+|    8 |            2.0 |  51.20 |  51.40 |  52.32 |  49.01 |  73.02 |  73.30 |  75.00 |  71.99 |  1.47 |
+|   16 |            2.0 |  51.75 |  52.24 |  56.36 |  51.27 |  83.99 |  84.57 |  86.69 |  83.24 |  1.62 |
+|    1 |            7.0 |  48.13 |  48.53 |  50.95 |  46.78 |  48.52 |  48.75 |  50.89 |  48.01 |  1.03 |
+|    2 |            7.0 |  49.52 |  50.10 |  52.35 |  48.00 |  65.27 |  65.41 |  66.59 |  64.79 |  1.35 |
+|    4 |            7.0 |  51.75 |  52.01 |  54.39 |  50.38 |  93.75 |  94.77 |  97.04 |  92.27 |  1.83 |
+|    8 |            7.0 |  54.80 |  56.27 |  66.23 |  52.95 | 130.65 | 131.09 | 132.91 | 129.82 |  2.45 |
+|   16 |            7.0 |  73.02 |  73.42 |  75.83 |  71.96 | 157.53 | 158.20 | 160.73 | 155.51 |  2.16 |
+|    1 |           16.7 |  48.10 |  48.52 |  52.71 |  47.20 |  73.34 |  73.56 |  74.19 |  72.69 |  1.54 |
+|    2 |           16.7 |  64.21 |  64.52 |  65.56 |  56.06 | 129.48 | 129.97 | 131.78 | 126.36 |  2.25 |
+|    4 |           16.7 |  60.38 |  61.03 |  63.18 |  58.87 | 183.33 | 183.85 | 185.53 | 181.90 |  3.09 |
+|    8 |           16.7 |  85.88 |  86.34 |  87.70 |  84.46 | 227.42 | 228.21 | 229.63 | 225.71 |  2.67 |
+|   16 |           16.7 | 135.62 | 136.40 | 137.69 | 131.58 | 276.90 | 277.59 | 281.16 | 275.08 |  2.09 |
+
+To achieve these same results, follow the [Quick Start Guide](#quick-start-guide) outlined above.
+
+##### Inference performance: NVIDIA DGX-1 (1x V100 32GB)
+|  |  | FP16 Latency (ms) Percentiles |  |  |  | FP32 Latency (ms) Percentiles |  |  |  | FP16/FP32 speed up |
+|-----:|---------------:|-------:|-------:|-------:|-------:|-------:|-------:|-------:|-------:|------:|
+|   BS |   Duration (s) |    90% |    95% |    99% |    Avg |    90% |    95% |    99% |    Avg |   Avg |
+|    1 |            2.0 |  52.74 |  53.01 |  54.40 |  51.47 |  55.97 |  56.22 |  57.93 |  54.93 |  1.07 |
+|    2 |            2.0 |  51.77 |  52.15 |  54.69 |  50.98 |  56.58 |  56.87 |  58.88 |  55.35 |  1.09 |
+|    4 |            2.0 |  51.41 |  51.76 |  53.47 |  50.55 |  61.56 |  61.87 |  63.81 |  60.74 |  1.20 |
+|    8 |            2.0 |  51.83 |  52.15 |  54.08 |  50.85 |  80.20 |  80.69 |  81.67 |  77.69 |  1.53 |
+|   16 |            2.0 |  70.48 |  70.96 |  72.11 |  62.98 |  93.00 |  93.44 |  94.17 |  89.05 |  1.41 |
+|    1 |            7.0 |  49.77 |  50.21 |  51.88 |  48.73 |  52.74 |  52.99 |  54.54 |  51.67 |  1.06 |
+|    2 |            7.0 |  51.12 |  51.47 |  52.84 |  49.98 |  65.33 |  65.63 |  67.07 |  64.64 |  1.29 |
+|    4 |            7.0 |  53.13 |  53.56 |  55.68 |  52.15 |  93.54 |  93.85 |  94.72 |  92.76 |  1.78 |
+|    8 |            7.0 |  57.67 |  58.07 |  59.89 |  56.41 | 133.93 | 134.18 | 134.88 | 133.15 |  2.36 |
+|   16 |            7.0 |  76.09 |  76.48 |  79.13 |  75.27 | 162.35 | 162.77 | 164.63 | 161.30 |  2.14 |
+|    1 |           16.7 |  54.78 |  55.29 |  56.83 |  52.51 |  75.37 |  76.27 |  78.05 |  74.32 |  1.42 |
+|    2 |           16.7 |  56.80 |  57.20 |  59.01 |  55.49 | 130.60 | 131.36 | 132.93 | 128.55 |  2.32 |
+|    4 |           16.7 |  64.19 |  64.84 |  66.47 |  62.87 | 188.09 | 188.76 | 190.07 | 185.76 |  2.95 |
+|    8 |           16.7 |  87.46 |  87.86 |  89.99 |  86.47 | 232.33 | 232.89 | 234.43 | 230.44 |  2.67 |
+|   16 |           16.7 | 136.02 | 136.52 | 139.44 | 134.78 | 283.87 | 284.59 | 286.70 | 282.01 |  2.09 |
+
+To achieve these same results, follow the [Quick Start Guide](#quick-start-guide) outlined above.
+
+##### Inference performance: NVIDIA DGX-2 (1x V100 32GB)
+|  |  | FP16 Latency (ms) Percentiles |  |  |  | FP32 Latency (ms) Percentiles |  |  |  | FP16/FP32 speed up |
+|-----:|---------------:|-------:|-------:|-------:|-------:|-------:|-------:|-------:|-------:|------:|
+|   BS |   Duration (s) |    90% |    95% |    99% |    Avg |    90% |    95% |    99% |    Avg |   Avg |
+|    1 |            2.0 |  35.88 |  36.12 |  39.80 |  35.20 |  42.95 |  43.67 |  46.65 |  42.23 |  1.20 |
+|    2 |            2.0 |  36.36 |  36.57 |  40.97 |  35.60 |  41.83 |  42.21 |  45.60 |  40.97 |  1.15 |
+|    4 |            2.0 |  36.69 |  36.89 |  41.25 |  36.05 |  48.35 |  48.52 |  52.35 |  47.80 |  1.33 |
+|    8 |            2.0 |  37.49 |  37.70 |  41.37 |  36.88 |  65.41 |  65.64 |  66.50 |  64.96 |  1.76 |
+|   16 |            2.0 |  41.35 |  41.79 |  45.58 |  40.91 |  77.22 |  77.51 |  79.48 |  76.54 |  1.87 |
+|    1 |            7.0 |  36.07 |  36.55 |  40.31 |  35.62 |  39.52 |  39.84 |  43.07 |  38.93 |  1.09 |
+|    2 |            7.0 |  37.42 |  37.66 |  41.36 |  36.79 |  55.94 |  56.19 |  58.33 |  55.60 |  1.51 |
+|    4 |            7.0 |  38.51 |  38.95 |  42.55 |  37.98 |  86.62 |  87.08 |  87.50 |  86.20 |  2.27 |
+|    8 |            7.0 |  42.82 |  43.00 |  47.11 |  42.55 | 122.05 | 122.29 | 122.70 | 121.59 |  2.86 |
+|   16 |            7.0 |  67.74 |  67.92 |  69.05 |  65.69 | 149.92 | 150.16 | 151.03 | 149.49 |  2.28 |
+|    1 |           16.7 |  39.28 |  39.78 |  43.34 |  38.35 |  66.73 |  67.16 |  69.80 |  66.01 |  1.72 |
+|    2 |           16.7 |  43.05 |  43.42 |  47.18 |  42.43 | 120.04 | 121.12 | 123.32 | 118.14 |  2.78 |
+|    4 |           16.7 |  52.18 |  52.49 |  56.11 |  51.63 | 176.09 | 176.51 | 178.70 | 174.60 |  3.38 |
+|    8 |           16.7 |  78.55 |  78.79 |  81.66 |  78.04 | 216.19 | 216.68 | 217.63 | 214.48 |  2.75 |
+|   16 |           16.7 | 125.57 | 125.92 | 128.78 | 124.33 | 264.11 | 264.49 | 266.14 | 262.80 |  2.11 |
+
+To achieve these same results, follow the [Quick Start Guide](#quick-start-guide) outlined above.
+
+##### Inference performance: NVIDIA T4
+|  |  | FP16 Latency (ms) Percentiles |  |  |  | FP32 Latency (ms) Percentiles |  |  |  | FP16/FP32 speed up |
+|-----:|---------------:|-------:|-------:|-------:|-------:|-------:|-------:|-------:|-------:|------:|
+|   BS |   Duration (s) |    90% |    95% |    99% |    Avg |    90% |    95% |    99% |    Avg |   Avg |
+|    1 |            2.0 |  43.62 |  46.95 |  50.46 |  37.23 |  51.31 |  52.37 |  56.21 |  49.77 |  1.34 |
+|    2 |            2.0 |  49.09 |  50.46 |  53.11 |  40.61 |  81.85 |  82.22 |  83.94 |  80.81 |  1.99 |
+|    4 |            2.0 |  47.71 |  51.14 |  55.09 |  41.29 | 112.56 | 115.13 | 118.56 | 111.60 |  2.70 |
+|    8 |            2.0 |  51.37 |  53.11 |  55.48 |  45.94 | 198.95 | 199.48 | 200.28 | 197.22 |  4.29 |
+|   16 |            2.0 |  63.59 |  64.30 |  66.90 |  61.77 | 221.75 | 222.07 | 223.22 | 220.09 |  3.56 |
+|    1 |            7.0 |  47.49 |  48.66 |  53.36 |  40.76 |  73.63 |  74.41 |  77.65 |  72.41 |  1.78 |
+|    2 |            7.0 |  48.63 |  50.01 |  58.35 |  43.44 | 114.66 | 115.28 | 117.63 | 112.41 |  2.59 |
+|    4 |            7.0 |  52.19 |  52.85 |  54.22 |  49.94 | 200.38 | 201.29 | 202.97 | 197.21 |  3.95 |
+|    8 |            7.0 |  84.90 |  85.56 |  87.52 |  83.41 | 404.00 | 404.72 | 405.70 | 400.25 |  4.80 |
+|   16 |            7.0 | 157.12 | 157.58 | 159.19 | 155.01 | 490.93 | 492.09 | 493.44 | 486.45 |  3.14 |
+|    1 |           16.7 |  50.57 |  51.57 |  57.58 |  46.27 | 150.39 | 151.84 | 153.54 | 147.31 |  3.18 |
+|    2 |           16.7 |  63.64 |  64.55 |  66.31 |  61.98 | 256.54 | 258.16 | 262.71 | 250.34 |  4.04 |
+|    4 |           16.7 | 140.44 | 141.06 | 142.00 | 138.14 | 519.59 | 521.41 | 523.86 | 512.74 |  3.71 |
+|    8 |           16.7 | 267.03 | 268.06 | 270.01 | 263.15 | 727.33 | 728.61 | 731.36 | 722.62 |  2.75 |
+|   16 |           16.7 | 362.40 | 364.02 | 367.80 | 358.75 | 867.92 | 869.19 | 871.46 | 860.37 |  2.40 |
+
+To achieve these same results, follow the [Quick Start Guide](#quick-start-guide) outlined above.
+
+## Release notes
+We're constantly refining and improving our performance on AI and HPC workloads even on the same hardware with frequent updates to our software stack. For our latest performance data please refer to these pages for AI and HPC benchmarks.
+
+### Changelog
+February 2021
+* Added DALI data-processing pipeline for on-the-fly data processing and augmentation on CPU or GPU
+* Revised training recipe: ~10% relative improvement in Word Error Rate (WER)
+* Updated Triton scripts for compatibility with Triton V2 API, updated Triton inference results
+* Refactored codebase
+* Updated performance results for the PyTorch 20.10-py3 NGC container
+
+June 2020
+* Updated performance tables to include A100 results
+
+December 2019
+* Inference support for TRT 6 with dynamic shapes
+* Inference support for TensorRT Inference Server with acoustic model backends in ONNX, PyTorch JIT, TensorRT
+* Jupyter notebook for inference with TensorRT Inference Server
+
+November 2019
+* Google Colab notebook for inference with native TensorRT
+
+September 2019
+* Inference support for TensorRT 6 with static shapes
+* Jupyter notebook for inference
+
+August 2019
+* Initial release
+
+### Known issues
+There are no known issues in this release.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.gitignore b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.gitignore
new file mode 100644
index 0000000000..d30f85c34f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.gitignore
@@ -0,0 +1,147 @@
+apex.egg-info
+dist
+build
+docs/build
+*~
+__pycache__
+.vscode
+
+# Copied from https://raw.githubusercontent.com/github/gitignore/master/Python.gitignore
+# Byte-compiled / optimized / DLL files
+__pycache__/
+*.py[cod]
+*$py.class
+
+# C extensions
+*.so
+
+# Distribution / packaging
+.Python
+build/
+develop-eggs/
+dist/
+downloads/
+eggs/
+.eggs/
+lib/
+lib64/
+parts/
+sdist/
+var/
+wheels/
+share/python-wheels/
+*.egg-info/
+.installed.cfg
+*.egg
+MANIFEST
+
+# PyInstaller
+#  Usually these files are written by a python script from a template
+#  before PyInstaller builds the exe, so as to inject date/other infos into it.
+*.manifest
+*.spec
+
+# Installer logs
+pip-log.txt
+pip-delete-this-directory.txt
+
+# Unit test / coverage reports
+htmlcov/
+.tox/
+.nox/
+.coverage
+.coverage.*
+.cache
+nosetests.xml
+coverage.xml
+*.cover
+*.py,cover
+.hypothesis/
+.pytest_cache/
+cover/
+
+# Translations
+*.mo
+*.pot
+
+# Django stuff:
+*.log
+local_settings.py
+db.sqlite3
+db.sqlite3-journal
+
+# Flask stuff:
+instance/
+.webassets-cache
+
+# Scrapy stuff:
+.scrapy
+
+# Sphinx documentation
+docs/_build/
+
+# PyBuilder
+.pybuilder/
+target/
+
+# Jupyter Notebook
+.ipynb_checkpoints
+
+# IPython
+profile_default/
+ipython_config.py
+
+# pyenv
+#   For a library or package, you might want to ignore these files since the code is
+#   intended to run in multiple environments; otherwise, check them in:
+# .python-version
+
+# pipenv
+#   According to pypa/pipenv#598, it is recommended to include Pipfile.lock in version control.
+#   However, in case of collaboration, if having platform-specific dependencies or dependencies
+#   having no cross-platform support, pipenv may install dependencies that don't work, or not
+#   install all needed dependencies.
+#Pipfile.lock
+
+# PEP 582; used by e.g. github.com/David-OConnor/pyflow
+__pypackages__/
+
+# Celery stuff
+celerybeat-schedule
+celerybeat.pid
+
+# SageMath parsed files
+*.sage.py
+
+# Environments
+.env
+.venv
+env/
+venv/
+ENV/
+env.bak/
+venv.bak/
+
+# Spyder project settings
+.spyderproject
+.spyproject
+
+# Rope project settings
+.ropeproject
+
+# mkdocs documentation
+/site
+
+# mypy
+.mypy_cache/
+.dmypy.json
+dmypy.json
+
+# Pyre type checker
+.pyre/
+
+# pytype static type analyzer
+.pytype/
+
+# Cython debug symbols
+cython_debug/
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.gitmodules b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.gitmodules
new file mode 100644
index 0000000000..6479428db0
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.gitmodules
@@ -0,0 +1,7 @@
+[submodule "apex/contrib/csrc/multihead_attn/cutlass"]
+	path = apex/contrib/csrc/multihead_attn/cutlass
+	url = https://github.com/NVIDIA/cutlass.git
+	branch = v1.2.0
+[submodule "apex/contrib/csrc/cudnn-frontend"]
+	path = apex/contrib/csrc/cudnn-frontend
+	url = https://github.com/NVIDIA/cudnn-frontend.git
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.nojekyll b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/.nojekyll
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/LICENSE b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/LICENSE
new file mode 100644
index 0000000000..3d1e9454ff
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/LICENSE
@@ -0,0 +1,11 @@
+All rights reserved.
+
+Redistribution and use in source and binary forms, with or without modification, are permitted provided that the following conditions are met:
+
+1. Redistributions of source code must retain the above copyright notice, this list of conditions and the following disclaimer.
+
+2. Redistributions in binary form must reproduce the above copyright notice, this list of conditions and the following disclaimer in the documentation and/or other materials provided with the distribution.
+
+3. Neither the name of the copyright holder nor the names of its contributors may be used to endorse or promote products derived from this software without specific prior written permission.
+
+THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT HOLDER OR CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/README.md
new file mode 100644
index 0000000000..a761def7b7
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/README.md
@@ -0,0 +1,146 @@
+# Introduction
+
+This repository holds NVIDIA-maintained utilities to streamline
+mixed precision and distributed training in Pytorch.
+Some of the code here will be included in upstream Pytorch eventually.
+The intention of Apex is to make up-to-date utilities available to
+users as quickly as possible.
+
+## Full API Documentation: [https://nvidia.github.io/apex](https://nvidia.github.io/apex)
+
+## [GTC 2019](https://github.com/mcarilli/mixed_precision_references/tree/master/GTC_2019) and [Pytorch DevCon 2019](https://github.com/mcarilli/mixed_precision_references/tree/master/Pytorch_Devcon_2019) Slides
+
+# Contents
+
+## 1. Amp:  Automatic Mixed Precision
+
+`apex.amp` is a tool to enable mixed precision training by changing only 3 lines of your script.
+Users can easily experiment with different pure and mixed precision training modes by supplying
+different flags to `amp.initialize`.
+
+[Webinar introducing Amp](https://info.nvidia.com/webinar-mixed-precision-with-pytorch-reg-page.html)
+(The flag `cast_batchnorm` has been renamed to `keep_batchnorm_fp32`).
+
+[API Documentation](https://nvidia.github.io/apex/amp.html)
+
+[Comprehensive Imagenet example](https://github.com/NVIDIA/apex/tree/master/examples/imagenet)
+
+[DCGAN example coming soon...](https://github.com/NVIDIA/apex/tree/master/examples/dcgan)
+
+[Moving to the new Amp API](https://nvidia.github.io/apex/amp.html#transition-guide-for-old-api-users) (for users of the deprecated "Amp" and "FP16_Optimizer" APIs)
+
+## 2. Distributed Training
+
+`apex.parallel.DistributedDataParallel` is a module wrapper, similar to
+`torch.nn.parallel.DistributedDataParallel`.  It enables convenient multiprocess distributed training,
+optimized for NVIDIA's NCCL communication library.
+
+[API Documentation](https://nvidia.github.io/apex/parallel.html)
+
+[Python Source](https://github.com/NVIDIA/apex/tree/master/apex/parallel)
+
+[Example/Walkthrough](https://github.com/NVIDIA/apex/tree/master/examples/simple/distributed)
+
+The [Imagenet example](https://github.com/NVIDIA/apex/tree/master/examples/imagenet)
+shows use of `apex.parallel.DistributedDataParallel` along with `apex.amp`.
+
+### Synchronized Batch Normalization
+
+`apex.parallel.SyncBatchNorm` extends `torch.nn.modules.batchnorm._BatchNorm` to
+support synchronized BN.
+It allreduces stats across processes during multiprocess (DistributedDataParallel) training.
+Synchronous BN has been used in cases where only a small
+local minibatch can fit on each GPU.
+Allreduced stats increase the effective batch size for the BN layer to the
+global batch size across all processes (which, technically, is the correct
+formulation).
+Synchronous BN has been observed to improve converged accuracy in some of our research models.
+
+### Checkpointing
+
+To properly save and load your `amp` training, we introduce the `amp.state_dict()`, which contains all `loss_scalers` and their corresponding unskipped steps,
+as well as `amp.load_state_dict()` to restore these attributes.
+
+In order to get bitwise accuracy, we recommend the following workflow:
+```python
+# Initialization
+opt_level = 'O1'
+model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
+
+# Train your model
+...
+with amp.scale_loss(loss, optimizer) as scaled_loss:
+    scaled_loss.backward()
+...
+
+# Save checkpoint
+checkpoint = {
+    'model': model.state_dict(),
+    'optimizer': optimizer.state_dict(),
+    'amp': amp.state_dict()
+}
+torch.save(checkpoint, 'amp_checkpoint.pt')
+...
+
+# Restore
+model = ...
+optimizer = ...
+checkpoint = torch.load('amp_checkpoint.pt')
+
+model, optimizer = amp.initialize(model, optimizer, opt_level=opt_level)
+model.load_state_dict(checkpoint['model'])
+optimizer.load_state_dict(checkpoint['optimizer'])
+amp.load_state_dict(checkpoint['amp'])
+
+# Continue training
+...
+```
+
+Note that we recommend restoring the model using the same `opt_level`. Also note that we recommend calling the `load_state_dict` methods after `amp.initialize`.
+
+# Requirements
+
+Python 3
+
+CUDA 9 or newer
+
+PyTorch 0.4 or newer.  The CUDA and C++ extensions require pytorch 1.0 or newer.
+
+We recommend the latest stable release, obtainable from
+[https://pytorch.org/](https://pytorch.org/).  We also test against the latest master branch, obtainable from [https://github.com/pytorch/pytorch](https://github.com/pytorch/pytorch).
+
+It's often convenient to use Apex in Docker containers.  Compatible options include:
+* [NVIDIA Pytorch containers from NGC](https://ngc.nvidia.com/catalog/containers/nvidia%2Fpytorch), which come with Apex preinstalled.  To use the latest Amp API, you may need to `pip uninstall apex` then reinstall Apex using the **Quick Start** commands below.
+* [official Pytorch -devel Dockerfiles](https://hub.docker.com/r/pytorch/pytorch/tags), e.g. `docker pull pytorch/pytorch:nightly-devel-cuda10.0-cudnn7`, in which you can install Apex using the **Quick Start** commands.
+
+See the [Docker example folder](https://github.com/NVIDIA/apex/tree/master/examples/docker) for details.
+
+# Quick Start
+
+### Linux
+
+For performance and full functionality, we recommend installing Apex with
+CUDA and C++ extensions via
+```
+git clone https://github.com/NVIDIA/apex
+cd apex
+pip install -v --disable-pip-version-check --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./
+```
+
+Apex also supports a Python-only build (required with Pytorch 0.4) via
+```
+pip install -v --disable-pip-version-check --no-cache-dir ./
+```
+A Python-only build omits:
+- Fused kernels required to use `apex.optimizers.FusedAdam`.
+- Fused kernels required to use `apex.normalization.FusedLayerNorm`.
+- Fused kernels that improve the performance and numerical stability of `apex.parallel.SyncBatchNorm`.
+- Fused kernels that improve the performance of `apex.parallel.DistributedDataParallel` and `apex.amp`.
+`DistributedDataParallel`, `amp`, and `SyncBatchNorm` will still be usable, but they may be slower.
+
+Pyprof support has been moved to its own [dedicated repository](https://github.com/NVIDIA/PyProf).
+The codebase is deprecated in Apex and will be removed soon.
+
+### Windows support
+Windows support is experimental, and Linux is recommended.  `pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .` may work if you were able to build Pytorch from source
+on your system.  `pip install -v --no-cache-dir .` (without CUDA/C++ extensions) is more likely to work.  If you installed Pytorch in a Conda environment, make sure to install Apex in that same environment.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/README.md
new file mode 100644
index 0000000000..9e86fd8fc1
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/README.md
@@ -0,0 +1 @@
+Under construction...
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/RNNBackend.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/RNNBackend.py
new file mode 100644
index 0000000000..b9d4937efa
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/RNNBackend.py
@@ -0,0 +1,365 @@
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+
+import torch.nn.functional as F
+
+import math
+
+
+def is_iterable(maybe_iterable):
+    return isinstance(maybe_iterable, list) or isinstance(maybe_iterable, tuple)
+
+
+def flatten_list(tens_list):
+    """
+    flatten_list
+    """
+    if not is_iterable(tens_list):
+        return tens_list
+    
+    return torch.cat(tens_list, dim=0).view(len(tens_list), *tens_list[0].size() )
+
+    
+#These modules always assumes batch_first
+class bidirectionalRNN(nn.Module):
+    """
+    bidirectionalRNN
+    """
+    def __init__(self, inputRNN, num_layers=1, dropout = 0):
+        super(bidirectionalRNN, self).__init__()
+        self.dropout = dropout
+        self.fwd = stackedRNN(inputRNN, num_layers=num_layers, dropout = dropout)
+        self.bckwrd = stackedRNN(inputRNN.new_like(), num_layers=num_layers, dropout = dropout)
+        self.rnns = nn.ModuleList([self.fwd, self.bckwrd])
+        
+    #collect hidden option will return all hidden/cell states from entire RNN
+    def forward(self, input, collect_hidden=False):
+        """
+        forward()
+        """
+        seq_len = input.size(0)
+        bsz = input.size(1)
+
+        fwd_out, fwd_hiddens = list(self.fwd(input, collect_hidden = collect_hidden))
+        bckwrd_out, bckwrd_hiddens = list(self.bckwrd(input, reverse=True, collect_hidden = collect_hidden))
+        
+        output = torch.cat( [fwd_out, bckwrd_out], -1 )
+        hiddens = tuple( torch.cat(hidden, -1) for hidden in zip( fwd_hiddens, bckwrd_hiddens) )
+
+        return output, hiddens
+
+    def reset_parameters(self):
+        """
+        reset_parameters()
+        """
+        for rnn in self.rnns:
+            rnn.reset_parameters()
+        
+    def init_hidden(self, bsz):
+        """
+        init_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.init_hidden(bsz)
+
+    def detach_hidden(self):
+        """
+        detach_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.detachHidden()
+        
+    def reset_hidden(self, bsz):
+        """
+        reset_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.reset_hidden(bsz)
+
+    def init_inference(self, bsz):    
+        """
+        init_inference()
+        """
+        for rnn in self.rnns:
+            rnn.init_inference(bsz)
+
+   
+#assumes hidden_state[0] of inputRNN is output hidden state
+#constructor either takes an RNNCell or list of RNN layers
+class stackedRNN(nn.Module):        
+    """
+    stackedRNN
+    """
+    def __init__(self, inputRNN, num_layers=1, dropout=0):
+        super(stackedRNN, self).__init__()
+        
+        self.dropout = dropout
+        
+        if isinstance(inputRNN, RNNCell):
+            self.rnns = [inputRNN]
+            for i in range(num_layers-1):
+                self.rnns.append(inputRNN.new_like(inputRNN.output_size))
+        elif isinstance(inputRNN, list):
+            assert len(inputRNN) == num_layers, "RNN list length must be equal to num_layers"
+            self.rnns=inputRNN
+        else:
+            raise RuntimeError()
+        
+        self.nLayers = len(self.rnns)
+        
+        self.rnns = nn.ModuleList(self.rnns)
+
+
+    '''
+    Returns output as hidden_state[0] Tensor([sequence steps][batch size][features])
+    If collect hidden will also return Tuple(
+        [n_hidden_states][sequence steps] Tensor([layer][batch size][features])
+    )
+    If not collect hidden will also return Tuple(
+        [n_hidden_states] Tensor([layer][batch size][features])
+    '''
+    def forward(self, input, collect_hidden=False, reverse=False):
+        """
+        forward()
+        """
+        seq_len = input.size(0)
+        bsz = input.size(1)
+        inp_iter = reversed(range(seq_len)) if reverse else range(seq_len)
+
+        hidden_states = [[] for i in range(self.nLayers)]
+        outputs = []
+
+        for seq in inp_iter:
+            for layer in range(self.nLayers):
+
+                if layer == 0:
+                    prev_out = input[seq]
+                    
+                outs = self.rnns[layer](prev_out)
+
+                if collect_hidden:
+                    hidden_states[layer].append(outs)
+                elif seq == seq_len-1:
+                    hidden_states[layer].append(outs)
+                    
+                prev_out = outs[0]
+
+            outputs.append(prev_out)
+
+        if reverse:
+            outputs = list(reversed(outputs))
+        '''
+        At this point outputs is in format:
+        list( [seq_length] x Tensor([bsz][features]) )
+        need to convert it to:
+        list( Tensor([seq_length][bsz][features]) )
+        '''
+        output = flatten_list(outputs)
+
+        '''
+        hidden_states at this point is in format:
+        list( [layer][seq_length][hidden_states] x Tensor([bsz][features]) )
+        need to convert it to:
+          For not collect hidden:
+            list( [hidden_states] x Tensor([layer][bsz][features]) )
+          For collect hidden:
+            list( [hidden_states][seq_length] x Tensor([layer][bsz][features]) )
+        '''
+        if not collect_hidden:
+            seq_len = 1
+        n_hid = self.rnns[0].n_hidden_states
+        new_hidden = [ [ [ None for k in range(self.nLayers)] for j in range(seq_len) ] for i in range(n_hid) ]
+
+
+        for i in range(n_hid):
+            for j in range(seq_len):
+                for k in range(self.nLayers):
+                    new_hidden[i][j][k] = hidden_states[k][j][i]
+
+        hidden_states = new_hidden
+        #Now in format list( [hidden_states][seq_length][layer] x Tensor([bsz][features]) )
+        #Reverse seq_length if reverse
+        if reverse:
+            hidden_states = list( list(reversed(list(entry))) for entry in hidden_states)
+
+        #flatten layer dimension into tensor
+        hiddens = list( list(
+            flatten_list(seq) for seq in hidden )
+                        for hidden in hidden_states )
+        
+        #Now in format list( [hidden_states][seq_length] x Tensor([layer][bsz][features]) )
+        #Remove seq_length dimension if not collect_hidden
+        if not collect_hidden:
+            hidden_states = list( entry[0] for entry in hidden_states)
+        return output, hidden_states
+    
+    def reset_parameters(self):
+        """
+        reset_parameters()
+        """
+        for rnn in self.rnns:
+            rnn.reset_parameters()
+        
+    def init_hidden(self, bsz):
+        """
+        init_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.init_hidden(bsz)
+
+    def detach_hidden(self):
+        """
+        detach_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.detach_hidden()
+        
+    def reset_hidden(self, bsz):
+        """
+        reset_hidden()
+        """
+        for rnn in self.rnns:
+            rnn.reset_hidden(bsz)
+
+    def init_inference(self, bsz):    
+        """ 
+        init_inference()
+        """
+        for rnn in self.rnns:
+            rnn.init_inference(bsz)
+
+class RNNCell(nn.Module):
+    """ 
+    RNNCell 
+    gate_multiplier is related to the architecture you're working with
+    For LSTM-like it will be 4 and GRU-like will be 3.
+    Always assumes input is NOT batch_first.
+    Output size that's not hidden size will use output projection
+    Hidden_states is number of hidden states that are needed for cell
+    if one will go directly to cell as tensor, if more will go as list
+    """
+    def __init__(self, gate_multiplier, input_size, hidden_size, cell, n_hidden_states = 2, bias = False, output_size = None):
+        super(RNNCell, self).__init__()
+
+        self.gate_multiplier = gate_multiplier
+        self.input_size = input_size
+        self.hidden_size = hidden_size
+        self.cell = cell
+        self.bias = bias
+        self.output_size = output_size
+        if output_size is None:
+            self.output_size = hidden_size
+
+        self.gate_size = gate_multiplier * self.hidden_size
+        self.n_hidden_states = n_hidden_states
+
+        self.w_ih = nn.Parameter(torch.Tensor(self.gate_size, self.input_size))
+        self.w_hh = nn.Parameter(torch.Tensor(self.gate_size, self.output_size))
+
+        #Check if there's recurrent projection
+        if(self.output_size != self.hidden_size):
+            self.w_ho = nn.Parameter(torch.Tensor(self.output_size, self.hidden_size))
+
+        self.b_ih = self.b_hh = None
+        if self.bias:
+            self.b_ih = nn.Parameter(torch.Tensor(self.gate_size))
+            self.b_hh = nn.Parameter(torch.Tensor(self.gate_size))
+            
+        #hidden states for forward
+        self.hidden = [ None for states in range(self.n_hidden_states)]
+
+        self.reset_parameters()
+
+    def new_like(self, new_input_size=None):
+        """
+        new_like()
+        """
+        if new_input_size is None:
+            new_input_size = self.input_size
+            
+        return type(self)(self.gate_multiplier,
+                       new_input_size,
+                       self.hidden_size,
+                       self.cell,
+                       self.n_hidden_states,
+                       self.bias,
+                       self.output_size)
+
+    
+    #Use xavier where we can (weights), otherwise use uniform (bias)
+    def reset_parameters(self, gain=1):
+        """
+        reset_parameters()
+        """
+        stdev = 1.0 / math.sqrt(self.hidden_size)
+        for param in self.parameters():
+            param.data.uniform_(-stdev, stdev)
+    '''
+    Xavier reset:
+    def reset_parameters(self, gain=1):
+        stdv = 1.0 / math.sqrt(self.gate_size)
+
+        for param in self.parameters():
+            if (param.dim() > 1):
+                torch.nn.init.xavier_normal(param, gain)
+            else:
+                param.data.uniform_(-stdv, stdv)
+    '''
+    def init_hidden(self, bsz):
+        """
+        init_hidden()
+        """
+        for param in self.parameters():
+            if param is not None:
+                a_param = param
+                break
+
+        for i, _ in enumerate(self.hidden):
+            if(self.hidden[i] is None or self.hidden[i].data.size()[0] != bsz):
+
+                if i==0:
+                    hidden_size = self.output_size
+                else:
+                    hidden_size = self.hidden_size
+
+                tens = a_param.data.new(bsz, hidden_size).zero_()
+                self.hidden[i] = Variable(tens, requires_grad=False)
+            
+        
+    def reset_hidden(self, bsz):
+        """
+        reset_hidden()
+        """
+        for i, _ in enumerate(self.hidden):
+            self.hidden[i] = None
+        self.init_hidden(bsz)
+
+    def detach_hidden(self):
+        """
+        detach_hidden()
+        """
+        for i, _ in enumerate(self.hidden):
+            if self.hidden[i] is None:
+                raise RuntimeError("Must initialize hidden state before you can detach it")
+        for i, _ in enumerate(self.hidden):
+            self.hidden[i] = self.hidden[i].detach()
+        
+    def forward(self, input):
+        """
+        forward()
+        if not inited or bsz has changed this will create hidden states
+        """
+        self.init_hidden(input.size()[0])
+
+        hidden_state = self.hidden[0] if self.n_hidden_states == 1 else self.hidden
+        self.hidden = self.cell(input, hidden_state, self.w_ih, self.w_hh, b_ih=self.b_ih, b_hh=self.b_hh)
+        if(self.n_hidden_states > 1):
+            self.hidden = list(self.hidden)
+        else:
+            self.hidden=[self.hidden]
+
+        if self.output_size != self.hidden_size:
+            self.hidden[0] = F.linear(self.hidden[0], self.w_ho)
+
+        return tuple(self.hidden)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/__init__.py
new file mode 100644
index 0000000000..d706746669
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/__init__.py
@@ -0,0 +1,3 @@
+from .models import LSTM, GRU, ReLU, Tanh, mLSTM
+
+__all__ = ['models']
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/cells.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/cells.py
new file mode 100644
index 0000000000..32b61a1be1
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/cells.py
@@ -0,0 +1,84 @@
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from .RNNBackend import RNNCell
+
+from torch.nn._functions.thnn import rnnFusedPointwise as fusedBackend
+
+import math 
+
+
+class mLSTMRNNCell(RNNCell):
+    """
+    mLSTMRNNCell
+    """
+
+    def __init__(self, input_size, hidden_size, bias = False, output_size = None):
+        gate_multiplier = 4
+        super(mLSTMRNNCell, self).__init__(gate_multiplier, input_size, hidden_size, mLSTMCell, n_hidden_states = 2, bias = bias, output_size = output_size)
+
+        self.w_mih = nn.Parameter(torch.Tensor(self.output_size, self.input_size))
+        self.w_mhh = nn.Parameter(torch.Tensor(self.output_size, self.output_size))
+
+        self.reset_parameters()
+
+    def forward(self, input):
+        """
+        mLSTMRNNCell.forward()
+        """
+        #if not inited or bsz has changed this will create hidden states
+        self.init_hidden(input.size()[0])
+
+        hidden_state = self.hidden[0] if self.n_hidden_states == 1 else self.hidden
+
+        self.hidden = list(
+                           self.cell(input, hidden_state, self.w_ih, self.w_hh, self.w_mih, self.w_mhh,
+                           b_ih=self.b_ih, b_hh=self.b_hh)
+        )
+        
+        if self.output_size != self.hidden_size:
+            self.hidden[0] = F.linear(self.hidden[0], self.w_ho)
+        return tuple(self.hidden)
+
+
+    def new_like(self, new_input_size=None):
+        if new_input_size is None:
+            new_input_size = self.input_size
+        
+        return type(self)(
+            new_input_size,
+            self.hidden_size,
+            self.bias,
+            self.output_size)
+
+def mLSTMCell(input, hidden, w_ih, w_hh, w_mih, w_mhh, b_ih=None, b_hh=None):
+    """
+    mLSTMCell
+    """
+
+    if input.is_cuda:
+        igates = F.linear(input, w_ih)
+        m = F.linear(input, w_mih) * F.linear(hidden[0], w_mhh)
+        hgates = F.linear(m, w_hh)
+
+        state = fusedBackend.LSTMFused.apply
+        return state(igates, hgates, hidden[1], b_ih, b_hh)
+
+    hx, cx = hidden
+    
+    m = F.linear(input, w_mih) * F.linear(hidden[0], w_mhh)
+    gates = F.linear(input, w_ih, b_ih) + F.linear(m, w_hh, b_hh)
+
+    ingate, forgetgate, cellgate, outgate = gates.chunk(4, 1)
+
+    ingate = F.sigmoid(ingate)
+    forgetgate = F.sigmoid(forgetgate)
+    cellgate = F.tanh(cellgate)
+    outgate = F.sigmoid(outgate)
+    
+    cy = (forgetgate * cx) + (ingate * cellgate)
+    hy = outgate * F.tanh(cy)
+    
+    return hy, cy
+                                                                            
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/models.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/models.py
new file mode 100644
index 0000000000..dd7adce047
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/RNN/models.py
@@ -0,0 +1,54 @@
+import torch
+
+from torch.nn._functions.rnn import LSTMCell, RNNReLUCell, RNNTanhCell, GRUCell
+
+from .RNNBackend import bidirectionalRNN, stackedRNN, RNNCell
+from .cells import mLSTMRNNCell, mLSTMCell
+
+def toRNNBackend(inputRNN, num_layers, bidirectional=False, dropout = 0):
+    """
+    :class:`toRNNBackend`
+    """
+
+    if bidirectional:
+        return bidirectionalRNN(inputRNN, num_layers, dropout = dropout)
+    else:
+        return stackedRNN(inputRNN, num_layers, dropout = dropout)
+
+
+def LSTM(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+    """
+    :class:`LSTM`
+    """
+    inputRNN = RNNCell(4, input_size, hidden_size, LSTMCell, 2, bias, output_size)
+    return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+def GRU(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+    """
+    :class:`GRU`
+    """
+    inputRNN = RNNCell(3, input_size, hidden_size, GRUCell, 1, bias, output_size)
+    return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+def ReLU(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+    """
+    :class:`ReLU`
+    """
+    inputRNN = RNNCell(1, input_size, hidden_size, RNNReLUCell, 1, bias, output_size)
+    return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+def Tanh(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+    """
+    :class:`Tanh`
+    """
+    inputRNN = RNNCell(1, input_size, hidden_size, RNNTanhCell, 1, bias, output_size)
+    return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+        
+def mLSTM(input_size, hidden_size, num_layers, bias=True, batch_first=False, dropout=0, bidirectional=False, output_size = None):
+    """
+    :class:`mLSTM`
+    """
+    inputRNN = mLSTMRNNCell(input_size, hidden_size, bias=bias, output_size=output_size)
+    return toRNNBackend(inputRNN, num_layers, bidirectional, dropout=dropout)
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/__init__.py
new file mode 100644
index 0000000000..42f8898f56
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/__init__.py
@@ -0,0 +1,20 @@
+# May help avoid undefined symbol errors https://pytorch.org/cppdocs/notes/faq.html#undefined-symbol-errors-from-pytorch-aten
+import torch
+
+if torch.distributed.is_available():
+    from . import parallel
+
+from . import amp
+from . import fp16_utils
+
+# For optimizers and normalization there is no Python fallback.
+# Absence of cuda backend is a hard error.
+# I would like the errors from importing fused_adam_cuda or fused_layer_norm_cuda
+# to be triggered lazily, because if someone has installed with --cpp_ext and --cuda_ext
+# so they expect those backends to be available, but for some reason they actually aren't
+# available (for example because they built improperly in a way that isn't revealed until
+# load time) the error message is timely and visible.
+from . import optimizers
+from . import normalization
+from . import pyprof
+from . import transformer
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/_autocast_utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/_autocast_utils.py
new file mode 100644
index 0000000000..94f8f91156
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/_autocast_utils.py
@@ -0,0 +1,17 @@
+from typing import Optional
+
+import torch
+
+
+def _get_current_dtype(dtype: Optional[torch.dtype] = None) -> torch.dtype:
+    if not torch.is_autocast_enabled():
+        return torch.float or dtype
+    else:
+        return torch.get_autocast_gpu_dtype()
+
+
+def _cast_if_autocast_enabled(*args):
+    if not torch.is_autocast_enabled():
+        return args
+    else:
+        return torch.cuda.amp.autocast_mode._cast(args, torch.get_autocast_gpu_dtype())
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/README.md
new file mode 100644
index 0000000000..a87b5010e3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/README.md
@@ -0,0 +1,72 @@
+# amp: Automatic Mixed Precision
+
+## Annotating User Functions
+
+Nearly all PyTorch user code needs nothing more than the two steps
+above to use amp. After all, custom layers are built out of simpler
+PyTorch components, and amp already can see those.
+
+However, any custom C++ or CUDA code is outside of amp's (default)
+view of things. For example, suppose I implemented a new recurrent
+cell called a "forgetful recurrent unit" that calls directly into a
+CUDA backend:
+
+```python
+from backend import FRUBackend
+
+def fru(input, hidden, weight, bias):
+    # call to CUDA code
+    FRUBackend(input, hidden, weight, bias)
+```
+
+In this case, it is possible to get a runtime type mismatch. For
+example, you might have `input` in fp16, and `weight` in fp32, and amp
+doesn't have the visibility to insert an appropriate cast.
+
+amp exposes two ways to handle "invisible" backend code: function
+annotations and explicit registration.
+
+#### Function annotation
+
+The first way to handle backend code is a set of function annotations:
+
+- `@amp.half_function`
+- `@amp.float_function`
+- `@amp.promote_function`
+
+These correspond to:
+
+- Cast all arguments to fp16
+- Cast all argumnets fo fp32
+- If there are any type mismatches, cast everything to the widest type
+
+In our example, we believe that the FRU unit is fp16-safe and will get
+performance gains from casting its arguments to fp16, so we write:
+
+```python
+@amp.half_function
+def fru(input, hidden, weight, bias):
+    #...
+```
+
+#### Explicit registration
+
+The other way to handle backend code is with explicit function
+registration:
+
+- `amp.register_half_function(module, function_name)`
+- `amp.register_float_function(module, function_name)`
+- `amp.register_promote_function(module, function_name)`
+
+When using this API, `module` is the containing class or module for
+the function, and `function_name` is the _string_ name of the
+function. Note that the function must be registered before the call to
+`amp.initalize()`.
+
+For our FRU unit, we can register the backend function directly:
+
+```python
+import backend
+
+amp.register_half_function(backend, 'FRUBackend')
+```
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/__init__.py
new file mode 100644
index 0000000000..34d080a69e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/__init__.py
@@ -0,0 +1,5 @@
+from .amp import init, half_function, float_function, promote_function,\
+    register_half_function, register_float_function, register_promote_function
+from .handle import scale_loss, disable_casts
+from .frontend import initialize, state_dict, load_state_dict
+from ._amp_state import master_params, _amp_state
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/__version__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/__version__.py
new file mode 100644
index 0000000000..3a83701b29
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/__version__.py
@@ -0,0 +1,2 @@
+VERSION = (0, 1, 0)
+__version__ = '.'.join(map(str, VERSION))
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_amp_state.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_amp_state.py
new file mode 100644
index 0000000000..1ac9d31160
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_amp_state.py
@@ -0,0 +1,69 @@
+# This is a "header object" that allows different amp modules to communicate.
+# I'm a C++ guy, not a python guy.  I decided this approach because it seemed most C++-like.
+# But apparently it's ok:
+# http://effbot.org/pyfaq/how-do-i-share-global-variables-across-modules.htm
+import os
+import torch
+
+TORCH_MAJOR = int(torch.__version__.split('.')[0])
+TORCH_MINOR = int(torch.__version__.split('.')[1])
+
+
+if TORCH_MAJOR == 1 and TORCH_MINOR < 8:
+    from torch._six import container_abcs
+else:
+    import collections.abc as container_abcs
+
+
+class AmpState(object):
+    def __init__(self):
+        self.hard_override=False
+        self.allow_incoming_model_not_fp32 = False
+        self.verbosity=1
+
+
+# Attribute stash.  Could also just stash things as global module attributes.
+_amp_state = AmpState()
+
+
+def warn_or_err(msg):
+    if _amp_state.hard_override:
+        print("Warning:  " + msg)
+    else:
+        raise RuntimeError(msg)
+        # I'm not sure if allowing hard_override is a good idea.
+        # + "  If you're sure you know what you're doing, supply " +
+        #                    "hard_override=True to amp.initialize.")
+
+
+def maybe_print(msg, rank0=False):
+    distributed = torch.distributed.is_available() and \
+        torch.distributed.is_initialized() and \
+        torch.distributed.get_world_size() > 1
+    if _amp_state.verbosity > 0:
+        if rank0:
+            if distributed:
+                if torch.distributed.get_rank() == 0:
+                    print(msg)
+            else:
+                print(msg)
+        else:
+            print(msg)
+
+
+# def iter_params(param_groups):
+#     for group in param_groups:
+#         for p in group['params']:
+#             yield p
+
+
+def master_params(optimizer):
+    """
+    Generator expression that iterates over the params owned by ``optimizer``.
+
+    Args:
+        optimizer: An optimizer previously returned from ``amp.initialize``.
+    """
+    for group in optimizer.param_groups:
+        for p in group['params']:
+            yield p
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_initialize.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_initialize.py
new file mode 100644
index 0000000000..28c5bbbdfe
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_initialize.py
@@ -0,0 +1,263 @@
+import torch
+from torch._six import string_classes
+import functools
+import numpy as np
+import sys
+from types import MethodType
+import warnings
+from ._amp_state import _amp_state, warn_or_err, container_abcs
+from .handle import disable_casts
+from .scaler import LossScaler
+from ._process_optimizer import _process_optimizer
+from apex.fp16_utils import convert_network
+from ..fp16_utils import FP16_Optimizer as FP16_Optimizer_general
+from ..contrib.optimizers import FP16_Optimizer as FP16_Optimizer_for_fused
+
+if torch.distributed.is_available():
+    from ..parallel import DistributedDataParallel as apex_DDP
+    from ..parallel.LARC import LARC
+
+
+def to_type(dtype, t):
+    if isinstance(t, torch.Tensor):
+        if not t.is_cuda:
+            # This should not be a hard error, since it may be legitimate.
+            warnings.warn("An input tensor was not cuda.")
+        # GANs require this.
+        # if t.requires_grad:
+        #     warn_or_err("input data requires grad.  Since input data is not a model parameter,\n"
+        #         "its gradients will not be properly allreduced by DDP.")
+        if t.is_floating_point():
+            return t.to(dtype)
+        return t
+    else:
+        # Trust the user's custom batch type, that's all I can do here.
+        return t.to(dtype)
+
+
+# Modified from torch.optim.optimizer.py.  This is a bit more general than casted_args in utils.py.
+def applier(value, fn):
+    if isinstance(value, torch.Tensor):
+        return fn(value)
+    elif isinstance(value, string_classes):
+        return value
+    elif isinstance(value, np.ndarray):
+        return value
+    elif hasattr(value, "to"): # Allow handling of custom batch classes
+        return fn(value)
+    elif isinstance(value, container_abcs.Mapping):
+        return {applier(k, fn) : applier(v, fn) for k, v in value.items()}
+    elif isinstance(value, container_abcs.Iterable):
+        return type(value)(applier(v, fn) for v in value)
+    else:
+        # Do I want this to fire off even if someone chooses to pass something ordinary like
+        # an int or float?  May be more annoying than it's worth.
+        # print("Warning:  unrecognized type in applier.  If your input data is a custom class, "
+        #     "provide it with a .to(dtype) method which converts its floating-point Tensors to dtype. "
+        #     "Amp will check for your custom to() and invoke it to cast the batch's "
+        #     "floating-point Tensors to the appropriate type. "
+        #     "Also, if your data is a custom class, it is your responsibility to ensure that "
+        #     "any Tensors you want to be cuda are already cuda."
+        return value
+
+
+def check_models(models):
+    for model in models:
+        parallel_type = None
+        if isinstance(model, torch.nn.parallel.DistributedDataParallel):
+            parallel_type = "torch.nn.parallel.DistributedDataParallel"
+        if ('apex_DDP' in sys.modules) and isinstance(model, apex_DDP):
+            parallel_type = "apex.parallel.DistributedDataParallel"
+        if isinstance(model, torch.nn.parallel.DataParallel):
+            parallel_type = "torch.nn.parallel.DataParallel"
+        if parallel_type is not None:
+            raise RuntimeError("Incoming model is an instance of {}. ".format(parallel_type) +
+                "Parallel wrappers should only be applied to the model(s) AFTER \n"
+                "the model(s) have been returned from amp.initialize.")
+
+
+def check_params_fp32(models):
+    for model in models:
+        for name, param in model.named_parameters():
+            if param.is_floating_point():
+                if 'Half' in param.type():
+                    warn_or_err("Found param {} with type {}, expected torch.cuda.FloatTensor.\n"
+                        "When using amp.initialize, you do not need to call .half() on your model\n"
+                        "before passing it, no matter what optimization level you choose.".format(
+                        name, param.type()))
+                elif not param.is_cuda:
+                    warn_or_err("Found param {} with type {}, expected torch.cuda.FloatTensor.\n"
+                        "When using amp.initialize, you need to provide a model with parameters\n"
+                        "located on a CUDA device before passing it no matter what optimization level\n"
+                        "you chose. Use model.to('cuda') to use the default device.".format(
+                        name, param.type()))
+
+        # Backward compatibility for PyTorch 0.4
+        if hasattr(model, 'named_buffers'):
+            buf_iter = model.named_buffers()
+        else:
+            buf_iter = model._buffers
+        for obj in buf_iter:
+            if type(obj)==tuple:
+                name, buf = obj
+            else:
+                name, buf = obj, buf_iter[obj]
+            if buf.is_floating_point():
+                if 'Half' in buf.type():
+                    warn_or_err("Found buffer {} with type {}, expected torch.cuda.FloatTensor.\n"
+                        "When using amp.initialize, you do not need to call .half() on your model\n"
+                        "before passing it, no matter what optimization level you choose.".format(
+                        name, buf.type()))
+                elif not buf.is_cuda:
+                    warn_or_err("Found buffer {} with type {}, expected torch.cuda.FloatTensor.\n"
+                        "When using amp.initialize, you need to provide a model with buffers\n"
+                        "located on a CUDA device before passing it no matter what optimization level\n"
+                        "you chose. Use model.to('cuda') to use the default device.".format(
+                        name, buf.type()))
+
+
+def check_optimizers(optimizers):
+    for optim in optimizers:
+        bad_optim_type = None
+        if isinstance(optim, FP16_Optimizer_general):
+            bad_optim_type = "apex.fp16_utils.FP16_Optimizer"
+        if isinstance(optim, FP16_Optimizer_for_fused):
+            bad_optim_type = "apex.optimizers.FP16_Optimizer"
+        if bad_optim_type is not None:
+            raise RuntimeError("An incoming optimizer is an instance of {}. ".format(bad_optim_type) +
+                               "The optimizer(s) passed to amp.initialize() must be bare \n"
+                               "instances of either ordinary Pytorch optimizers, or Apex fused \n"
+                               "optimizers.\n")
+
+
+class O2StateDictHook(object):
+    def __init__(self, fn):
+        self.fn = fn
+
+    def __call__(self, module, state_dict, prefix, local_metadata):
+        for key in state_dict:
+            param = state_dict[key]
+            if 'Half' in param.type():
+                param = param.to(torch.float32)
+                state_dict[key] = param
+
+
+def _initialize(models, optimizers, properties, num_losses=1, cast_model_outputs=None):
+    from .amp import init as amp_init
+
+    optimizers_was_list = False
+    if isinstance(optimizers, torch.optim.Optimizer) or ('LARC' in globals() and isinstance(optimizers, LARC)):
+        optimizers = [optimizers]
+    elif optimizers is None:
+        optimizers = []
+    elif isinstance(optimizers, list):
+        optimizers_was_list = True
+        check_optimizers(optimizers)
+    else:
+        check_optimizers([optimizers])
+        raise TypeError("optimizers must be either a single optimizer or a list of optimizers.")
+
+    if isinstance(models, torch.nn.Module):
+        models_was_list = False
+        models = [models]
+    elif isinstance(models, list):
+        models_was_list = True
+    else:
+        raise TypeError("models must be either a single model or a list of models.")
+
+    check_models(models)
+
+    if not _amp_state.allow_incoming_model_not_fp32:
+        check_params_fp32(models)
+
+    # In the future, when FP16_Optimizer can be deprecated and master weights can
+    # become an attribute, remember to stash master weights before casting the model.
+
+    if properties.cast_model_type:
+        if properties.keep_batchnorm_fp32:
+            for model in models:
+                convert_network(model, properties.cast_model_type)
+        else:
+            for model in models:
+                model.to(properties.cast_model_type)
+
+        input_caster = functools.partial(to_type, properties.cast_model_type)
+        if cast_model_outputs is not None:
+            output_caster = functools.partial(to_type, cast_model_outputs)
+        else:
+            output_caster = functools.partial(to_type, torch.float32)
+
+        for model in models:
+            # Patch the forward method to cast incoming data to the correct type, and
+            # outgoing data to float32, so "the user never needs to call .half()."
+            # I like writing things explicitly more than decorators.
+            def patch_forward(old_fwd):
+                def new_fwd(*args, **kwargs):
+                    output = old_fwd(*applier(args, input_caster),
+                                     **applier(kwargs, input_caster))
+                    return applier(output, output_caster)
+                return new_fwd
+
+            model.forward = patch_forward(model.forward)
+
+        # State dict trick to recast any preexisting per-param state tensors
+        for optimizer in optimizers:
+            optimizer.load_state_dict(optimizer.state_dict())
+
+        # patch model.state_dict() to return float32 params
+        for model in models:
+            for module in model.modules():
+                module._register_state_dict_hook(O2StateDictHook(functools.partial(to_type, torch.float32)))
+
+    elif cast_model_outputs is not None:
+        output_caster = functools.partial(to_type, cast_model_outputs)
+
+        for model in models:
+            def patch_forward(old_fwd):
+                def new_fwd(*args, **kwargs):
+                    output = old_fwd(*args, **kwargs)
+                    return applier(output, output_caster)
+                return new_fwd
+
+            model.forward = patch_forward(model.forward)
+
+    for i, optimizer in enumerate(optimizers):
+        optimizers[i] = _process_optimizer(optimizer, properties)
+
+    _amp_state.loss_scalers = []
+    for _ in range(num_losses):
+        _amp_state.loss_scalers.append(LossScaler(properties.loss_scale,
+                                                  min_loss_scale=_amp_state.min_loss_scale,
+                                                  max_loss_scale=_amp_state.max_loss_scale))
+
+    if properties.patch_torch_functions:
+        # handle is unused here. It's accessible later through a global value anyway.
+        handle = amp_init(loss_scale=properties.loss_scale, verbose=(_amp_state.verbosity == 2))
+        for optimizer in optimizers:
+            # Disable Amp casting for the optimizer step, because it should only be
+            # applied to FP32 master params anyway.
+            def patch_step(old_step):
+                def new_step(self, *args, **kwargs):
+                    with disable_casts():
+                        output = old_step(*args, **kwargs)
+                    return output
+                return new_step
+
+            optimizer.step = MethodType(patch_step(optimizer.step), optimizer)
+
+    if optimizers_was_list:
+        if models_was_list:
+            return models, optimizers
+        else:
+            return models[0], optimizers
+    else:
+        if models_was_list:
+            if len(optimizers) == 0:
+                return models
+            else:
+                return models, optimizers[0]
+        else:
+            if len(optimizers) == 0:
+                return models[0]
+            else:
+                return models[0], optimizers[0]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_process_optimizer.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_process_optimizer.py
new file mode 100644
index 0000000000..471289bba6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/_process_optimizer.py
@@ -0,0 +1,489 @@
+import types
+from ..fp16_utils import master_params_to_model_params
+from ..multi_tensor_apply import multi_tensor_applier
+from ._amp_state import maybe_print
+import torch
+from ..optimizers import FusedSGD
+
+
+class AmpOptimizerState(object):
+    def __init__(self):
+        pass
+
+
+def _master_params_to_model_params(self):
+    stash = self._amp_stash
+    if multi_tensor_applier.available:
+        if len(stash.all_fp16_params) > 0:
+            multi_tensor_applier(
+                stash.multi_tensor_scale,
+                stash.dummy_overflow_buf,
+                [stash.all_fp32_from_fp16_params, stash.all_fp16_params],
+                1.0)
+    else:
+        for fp16_group, fp32_from_fp16_group in zip(stash.fp16_groups, stash.fp32_from_fp16_groups):
+            master_params_to_model_params(fp16_group, fp32_from_fp16_group)
+
+
+def lazy_init_with_master_weights(self):
+        stash = self._amp_stash
+        stash.fp16_groups = []
+        stash.fp32_from_fp16_groups = []
+        stash.fp32_from_fp32_groups = []
+        for i, param_group in enumerate(self.param_groups):
+            # maybe_print("FP16_Optimizer processing param group {}:".format(i))
+            fp16_params_this_group = []
+            fp32_params_this_group = []
+            fp32_from_fp16_params_this_group = []
+            for i, param in enumerate(param_group['params']):
+                if param.requires_grad:
+                    if param.type() == 'torch.cuda.HalfTensor':
+                        # maybe_print("FP16_Optimizer received torch.cuda.HalfTensor with {}"
+                        #             .format(param.size()))
+                        fp16_params_this_group.append(param)
+                        master_param = param.detach().clone().float()
+                        master_param.requires_grad = True
+                        param_group['params'][i] = master_param
+                        fp32_from_fp16_params_this_group.append(master_param)
+                        # Reset existing state dict key to the new master param.
+                        # We still need to recast per-param state tensors, if any, to FP32.
+                        if param in self.state:
+                           self.state[master_param] = self.state.pop(param)
+                    elif param.type() == 'torch.cuda.FloatTensor':
+                        # maybe_print("FP16_Optimizer received torch.cuda.FloatTensor with {}"
+                        #             .format(param.size()))
+                        fp32_params_this_group.append(param)
+                        param_group['params'][i] = param
+                    else:
+                        raise TypeError("Optimizer's parameters must be either "
+                                        "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+                                        "Received {}".format(param.type()))
+
+            stash.fp16_groups.append(fp16_params_this_group)
+            stash.fp32_from_fp16_groups.append(fp32_from_fp16_params_this_group)
+            stash.fp32_from_fp32_groups.append(fp32_params_this_group)
+
+        stash.all_fp16_params = []
+        for group in stash.fp16_groups:
+            stash.all_fp16_params += group
+
+        stash.all_fp32_from_fp16_params = []
+        for group in stash.fp32_from_fp16_groups:
+            stash.all_fp32_from_fp16_params += group
+
+        stash.all_fp32_from_fp32_params = []
+        for group in stash.fp32_from_fp32_groups:
+            stash.all_fp32_from_fp32_params += group
+
+        # all_fp16_grad_stash is only needed for fused optimizers.
+        stash.all_fp16_grad_stash = [None for _ in stash.all_fp16_params]
+        # stash.all_fp32_from_fp16_grad_stash = [None for _ in stash.all_fp32_from_fp16_params]
+        stash.all_fp32_from_fp32_grad_stash = [None for _ in stash.all_fp32_from_fp32_params]
+
+        for param in stash.all_fp32_from_fp16_params:
+            param.grad = None
+
+        for param in stash.all_fp32_from_fp32_params:
+            param.grad = None
+
+        # Leverage state_dict() and load_state_dict() to recast preexisting per-param state tensors
+        self.load_state_dict(self.state_dict())
+
+
+def post_backward_models_are_masters(scaler, params, stashed_grads, scale_override=None):
+        grads_have_scale, stashed_have_scale, out_scale = scaler.loss_scale(), 1.0, 1.0
+
+        # not much to do if scale == 1.0 and static scaling
+        if scaler.loss_scale() == 1.0 and not scaler.dynamic:
+            # Clear the stash.
+            for i in range(len(stashed_grads)):
+                stashed_grads[i] = None
+            return
+        
+        if scale_override is not None:
+            grads_have_scale, stashed_have_scale, out_scale = scale_override
+
+        # This is a lot of python overhead...
+        grads_needing_unscale = []
+        grads_needing_unscale_with_stash = []
+        stashed = []
+        for param, stashed_grad in zip(params, stashed_grads):
+            if param.grad is None and stashed_grad is not None:
+                param.grad = stashed_grad
+            elif param.grad is not None and stashed_grad is None:
+                grads_needing_unscale.append(param.grad)
+            elif param.grad is not None and stashed_grad is not None:
+                grads_needing_unscale_with_stash.append(param.grad)
+                stashed.append(stashed_grad)
+            else: # param.grad is None and stashed_grad is None
+                continue
+
+        # unscale() implements grads*(1/scale), so "scale" should be grads_have_scale/out_scale.
+        if len(grads_needing_unscale) > 0:
+            scaler.unscale(
+                grads_needing_unscale,
+                grads_needing_unscale,
+                None, # unused_scale, currently present to avoid API breakage elsewhere
+                models_are_masters=True,
+                scale_override=grads_have_scale/out_scale)
+
+        if len(grads_needing_unscale_with_stash) > 0:
+            scaler.unscale_with_stashed(
+                grads_needing_unscale_with_stash,
+                stashed,
+                grads_needing_unscale_with_stash,
+                scale_override=(grads_have_scale, stashed_have_scale, out_scale))
+
+        # Clear the stash.
+        for i in range(len(stashed_grads)):
+            stashed_grads[i] = None
+
+
+def prepare_backward_with_master_weights(self):
+    stash = self._amp_stash
+
+    self._amp_lazy_init()
+
+    for i, param in enumerate(stash.all_fp16_params):
+        # Set up to leverage grad copy elision.
+        # This may behave differently from an unpatched optimizer if zero_grad is used and the param is unused.
+        param.grad = None
+
+    # for i, param in enumerate(stash.all_fp32_from_fp16_params):
+    #     stash.all_fp32_from_fp16_grad_stash[i] = param.grad
+
+    for i, param in enumerate(stash.all_fp32_from_fp32_params):
+        stash.all_fp32_from_fp32_grad_stash[i] = param.grad
+        # Set up to leverage grad copy elision:
+        param.grad = None
+
+
+def post_backward_with_master_weights(self, scaler):
+    stash = self._amp_stash
+
+    self._amp_lazy_init()
+
+    # This is a lot of python overhead...
+    fp16_grads_needing_unscale = []
+    new_fp32_grads = []
+    fp16_grads_needing_unscale_with_stash = []
+    preexisting_fp32_grads = []
+    for fp16_param, fp32_param in zip(stash.all_fp16_params,
+                                      stash.all_fp32_from_fp16_params):
+        if fp16_param.grad is None and fp32_param.grad is not None:
+            continue
+        elif fp16_param.grad is not None and fp32_param.grad is None:
+            fp32_param.grad = torch.empty_like(fp32_param)
+            fp16_grads_needing_unscale.append(fp16_param.grad)
+            new_fp32_grads.append(fp32_param.grad)
+        elif fp16_param.grad is not None and fp32_param.grad is not None:
+            fp16_grads_needing_unscale_with_stash.append(fp16_param.grad)
+            preexisting_fp32_grads.append(fp32_param.grad)
+        else: # fp16_param.grad is None and fp32_param.grad is None:
+            continue
+
+    if len(fp16_grads_needing_unscale) > 0:
+        scaler.unscale(
+            fp16_grads_needing_unscale,
+            new_fp32_grads,
+            scaler.loss_scale(),
+            models_are_masters=False)
+
+    if len(fp16_grads_needing_unscale_with_stash) > 0:
+        scaler.unscale_with_stashed(
+            fp16_grads_needing_unscale_with_stash,
+            preexisting_fp32_grads,
+            preexisting_fp32_grads)
+
+    # fp32 params can be treated as they would be in the "no_master_weights" case.
+    post_backward_models_are_masters(
+        scaler,
+        stash.all_fp32_from_fp32_params,
+        stash.all_fp32_from_fp32_grad_stash)
+
+
+def lazy_init_no_master_weights(self):
+    stash = self._amp_stash
+    stash.all_fp16_params = []
+    stash.all_fp32_params = []
+    for i, param_group in enumerate(self.param_groups):
+        for i, param in enumerate(param_group['params']):
+            if param.type() == 'torch.cuda.HalfTensor':
+                stash.all_fp16_params.append(param)
+            elif param.type() == 'torch.cuda.FloatTensor':
+                stash.all_fp32_params.append(param)
+            else:
+                raise TypeError("Optimizer's parameters must be either "
+                                "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+                                "Received {}".format(param.type()))
+
+    stash.all_fp16_grad_stash = [None for _ in stash.all_fp16_params]
+    stash.all_fp32_grad_stash = [None for _ in stash.all_fp32_params]
+
+
+def prepare_backward_no_master_weights(self):
+    stash = self._amp_stash
+
+    self._amp_lazy_init()
+
+    for i, param in enumerate(stash.all_fp16_params):
+        stash.all_fp16_grad_stash[i] = param.grad
+        # Set up to leverage grad copy elision:
+        param.grad = None
+
+    for i, param in enumerate(stash.all_fp32_params):
+        stash.all_fp32_grad_stash[i] = param.grad
+        # Set up to leverage grad copy elision:
+        param.grad = None
+
+
+def post_backward_no_master_weights(self, scaler):
+    stash = self._amp_stash
+
+    self._amp_lazy_init()
+
+    split_types = ((stash.all_fp16_params, stash.all_fp16_grad_stash),
+             (stash.all_fp32_params, stash.all_fp32_grad_stash))
+
+    for params, stashed_grads in split_types:
+        post_backward_models_are_masters(scaler, params, stashed_grads)
+
+
+#####################################################################################
+# FusedSGD versions
+#####################################################################################
+
+# FusedSGD never explicitly materializes the fp32 gradients for "fp32 from fp16" master params
+# outside the kernel, so we must accumulate directly into the model grads.
+def prepare_backward_with_master_weights_FusedSGD(self):
+    if self.materialize_master_grads:
+        prepare_backward_with_master_weights(self)
+    else:
+        stash = self._amp_stash
+
+        self._amp_lazy_init()
+
+        for i, param in enumerate(stash.all_fp16_params):
+            stash.all_fp16_grad_stash[i] = param.grad
+            # Set up to leverage grad copy elision:
+            param.grad = None
+
+        for i, param in enumerate(stash.all_fp32_from_fp32_params):
+            stash.all_fp32_from_fp32_grad_stash[i] = param.grad
+            # Set up to leverage grad copy elision:
+            param.grad = None
+
+
+def post_backward_with_master_weights_FusedSGD(self, scaler):
+    if self.materialize_master_grads:
+        post_backward_with_master_weights(self, scaler)
+    else:
+        stash = self._amp_stash
+
+        self._amp_lazy_init()
+
+        grads_have_scale = scaler.loss_scale()
+        stashed_have_scale = self.most_recent_scale
+        out_scale = grads_have_scale
+        if self.scale_set_by_backward:
+            out_scale = min(grads_have_scale, self.most_recent_scale)
+
+        split_types = ((stash.all_fp16_params, stash.all_fp16_grad_stash),
+                 (stash.all_fp32_from_fp32_params, stash.all_fp32_from_fp32_grad_stash))
+
+
+        # unscale_with_stashed() implements grads*1/scale + stashed_grads*1.
+        # stashed_grads are scaled by self.most_recent_scale.
+        for params, stashed_grads in split_types:
+            post_backward_models_are_masters(scaler, params, stashed_grads,
+                                             (grads_have_scale, stashed_have_scale, out_scale))
+
+        self.most_recent_scale = out_scale
+        self.scale_set_by_backward = True
+
+
+def prepare_backward_no_master_weights_FusedSGD(self):
+    prepare_backward_no_master_weights(self)
+
+
+def post_backward_no_master_weights_FusedSGD(self, scaler):
+    post_backward_no_master_weights(self, scaler)
+
+
+def _amp_lazy_init(self):
+    stash = self._amp_stash
+
+    if not stash.lazy_init_called:
+        self._lazy_init_maybe_master_weights()
+        stash.lazy_init_called = True
+
+
+def _process_optimizer(optimizer, properties):
+    if hasattr(optimizer, "_amp_stash"):
+        raise RuntimeError("A given optimizer should only be passed through amp.initialize once.")
+    else:
+        optimizer._amp_stash = AmpOptimizerState()
+
+    optimizer._amp_stash.lazy_init_called = False
+    optimizer._amp_stash.already_patched = False
+    optimizer._amp_stash.params_have_scaled_gradients = False
+
+    for name in ("_lazy_init_maybe_master_weights",
+                 "_master_params_to_model_params",
+                 "_prepare_amp_backward",
+                 "_post_amp_backward",
+                 "_amp_lazy_init"):
+        if hasattr(optimizer, name):
+            raise RuntimeError("Incoming optimizer already has {} defined.".format(name))
+
+    # TODO:  Centralize exposure and import error checking for the C backend.
+    if multi_tensor_applier.available:
+        import amp_C
+        optimizer._amp_stash.multi_tensor_scale = amp_C.multi_tensor_scale
+        optimizer._amp_stash.multi_tensor_l2norm = amp_C.multi_tensor_l2norm
+        optimizer._amp_stash.dummy_overflow_buf = torch.cuda.IntTensor([0]);
+
+    if properties.master_weights:
+        optimizer._lazy_init_maybe_master_weights = types.MethodType(
+            lazy_init_with_master_weights, optimizer)
+
+        optimizer._master_params_to_model_params = types.MethodType(
+            _master_params_to_model_params, optimizer)
+
+        old_step = optimizer.step
+        def new_step(self, closure=None):
+            if closure is not None:
+                raise RuntimeError("Currently, Amp does not support closure use with optimizers.")
+            retval = old_step()
+            if not isinstance(self, FusedSGD):
+                self._master_params_to_model_params()
+            # Clear the master grads that wouldn't be zeroed by model.zero_grad()
+            for param in self._amp_stash.all_fp32_from_fp16_params:
+                param.grad = None
+            return retval
+        optimizer.step = types.MethodType(new_step, optimizer)
+
+        old_zero_grad = optimizer.zero_grad
+        def new_zero_grad(self):
+            stash = self._amp_stash
+            self._amp_lazy_init()
+            # Zero the model grads.
+            for param in stash.all_fp16_params:
+                if param.grad is not None:
+                    param.grad.detach_()
+                    param.grad.zero_()
+            for param in stash.all_fp32_from_fp32_params:
+                if param.grad is not None:
+                    param.grad.detach_()
+                    param.grad.zero_()
+            # Clear the master grads that are independent of model grads
+            for param in self._amp_stash.all_fp32_from_fp16_params:
+                param.grad = None
+        optimizer.zero_grad = types.MethodType(new_zero_grad, optimizer)
+
+        if isinstance(optimizer, FusedSGD):
+            optimizer._prepare_amp_backward = types.MethodType(
+                prepare_backward_with_master_weights_FusedSGD, optimizer)
+            optimizer._post_amp_backward = types.MethodType(
+                post_backward_with_master_weights_FusedSGD, optimizer)
+        else:
+            optimizer._prepare_amp_backward = types.MethodType(
+                prepare_backward_with_master_weights, optimizer)
+            optimizer._post_amp_backward = types.MethodType(
+                post_backward_with_master_weights, optimizer)
+    else:
+        optimizer._lazy_init_maybe_master_weights = types.MethodType(
+            lazy_init_no_master_weights, optimizer)
+
+        if isinstance(optimizer, FusedSGD):
+            optimizer._prepare_amp_backward = types.MethodType(
+                prepare_backward_no_master_weights_FusedSGD, optimizer)
+            optimizer._post_amp_backward = types.MethodType(
+                post_backward_no_master_weights_FusedSGD, optimizer)
+        else:
+            optimizer._prepare_amp_backward = types.MethodType(
+                prepare_backward_no_master_weights, optimizer)
+            optimizer._post_amp_backward = types.MethodType(
+                post_backward_no_master_weights, optimizer)
+
+    optimizer._amp_lazy_init = types.MethodType(_amp_lazy_init, optimizer)
+
+    old_add_param_group = optimizer.add_param_group
+
+    def new_add_param_group(self, new_group):
+        stash = self._amp_stash
+
+        if not stash.lazy_init_called:
+            self._lazy_init_maybe_master_weights()
+            stash.lazy_init_called = True
+
+        assert isinstance(new_group, dict), "param group must be a dict"
+
+        new_params = new_group['params']
+        if isinstance(new_params, torch.Tensor):
+            new_group['params'] = [new_params]
+        elif isinstance(new_params, set):
+            raise TypeError('optimizer parameters need to be organized in ordered collections, but '
+                            'the ordering of tensors in sets will change between runs. Please use a list instead.')
+        else:
+            new_group['params'] = list(new_params)
+
+        if properties.master_weights:
+            # Mutate new_group in-place to use FP32 master params
+            fp16_params_this_group = []
+            fp32_params_this_group = []
+            fp32_from_fp16_params_this_group = []
+            for i, param in enumerate(new_group['params']):
+                if param.requires_grad:
+                    if param.type() == 'torch.cuda.HalfTensor':
+                        fp16_params_this_group.append(param)
+                        master_param = param.detach().clone().float()
+                        master_param.requires_grad = True
+                        new_group['params'][i] = master_param
+                        fp32_from_fp16_params_this_group.append(master_param)
+                    elif param.type() == 'torch.cuda.FloatTensor':
+                        fp32_params_this_group.append(param)
+                        new_group['params'][i] = param
+                    else:
+                        raise TypeError("Optimizer's parameters must be either "
+                                        "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+                                        "Received {}".format(param.type()))
+
+            stash.fp16_groups.append(fp16_params_this_group)
+            stash.fp32_from_fp16_groups.append(fp32_from_fp16_params_this_group)
+            stash.fp32_from_fp32_groups.append(fp32_params_this_group)
+
+            stash.all_fp16_params += fp16_params_this_group
+            stash.all_fp32_from_fp16_params += fp32_from_fp16_params_this_group
+            stash.all_fp32_from_fp32_params += fp32_params_this_group
+
+            # stash.all_fp32_from_fp16_grad_stash = [None for _ in stash.all_fp32_from_fp16_params]
+            stash.all_fp32_from_fp32_grad_stash += [None for _ in fp32_params_this_group]
+
+            # It should be ok to let params be added with existing .grad attributes.
+            # for param in fp16_params_this_group:
+            #     param.grad = None
+
+            # for param in fp32_from_fp16_params_this_group:
+            #     param.grad = None
+
+            # for param in stash.fp32_params_this_group:
+            #     param.grad = None
+        else:
+            for param in new_group['params']:
+                if param.type() == 'torch.cuda.HalfTensor':
+                    stash.all_fp16_params.append(param)
+                    stash.all_fp16_grad_stash.append(None)
+                elif param.type() == 'torch.cuda.FloatTensor':
+                    stash.all_fp32_params.append(param)
+                    stash.all_fp32_grad_stash.append(None)
+                else:
+                    raise TypeError("Optimizer's parameters must be either "
+                                    "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "
+                                    "Received {}".format(param.type()))
+
+        old_add_param_group(new_group)
+
+    optimizer.add_param_group = types.MethodType(new_add_param_group, optimizer)
+
+    return optimizer
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/amp.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/amp.py
new file mode 100644
index 0000000000..1eed72d07b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/amp.py
@@ -0,0 +1,177 @@
+from . import compat, rnn_compat, utils, wrap
+from .handle import AmpHandle, NoOpHandle
+from .lists import functional_overrides, torch_overrides, tensor_overrides
+from ._amp_state import _amp_state
+from .frontend import *
+
+import functools
+import itertools
+
+import torch
+
+
+_DECORATOR_HANDLE = None
+_USER_CAST_REGISTRY = set()
+_USER_PROMOTE_REGISTRY = set()
+
+
+def _decorator_helper(orig_fn, cast_fn, wrap_fn):
+    def wrapper(*args, **kwargs):
+        handle = _DECORATOR_HANDLE
+        if handle is None or not handle.is_active():
+            return orig_fn(*args, **kwargs)
+        inner_cast_fn = utils.verbosify(cast_fn, orig_fn.__name__,
+                                  handle.verbose)
+        return wrap_fn(orig_fn, inner_cast_fn, handle)(*args, **kwargs)
+    return wrapper
+
+
+# Decorator form
+def half_function(fn):
+    wrap_fn = functools.partial(wrap.make_cast_wrapper, try_caching=True)
+    return _decorator_helper(fn, utils.maybe_half, wrap_fn)
+
+
+def float_function(fn):
+    wrap_fn = functools.partial(wrap.make_cast_wrapper, try_caching=False)
+    return _decorator_helper(fn, utils.maybe_float, wrap_fn)
+
+
+def promote_function(fn):
+    wrap_fn = functools.partial(wrap.make_promote_wrapper)
+    return _decorator_helper(fn, utils.maybe_float, wrap_fn)
+
+
+# Registry form
+def register_half_function(module, name):
+    if not hasattr(module, name):
+        raise ValueError('No function named {} in module {}.'.format(
+            name, module))
+    _USER_CAST_REGISTRY.add((module, name, utils.maybe_half))
+
+
+def register_float_function(module, name):
+    if not hasattr(module, name):
+        raise ValueError('No function named {} in module {}.'.format(
+            name, module))
+    _USER_CAST_REGISTRY.add((module, name, utils.maybe_float))
+
+
+def register_promote_function(module, name):
+    if not hasattr(module, name):
+        raise ValueError('No function named {} in module {}.'.format(
+            name, module))
+    _USER_PROMOTE_REGISTRY.add((module, name))
+
+
+# Top-level function to insert _all_ the hooks.
+def init(enabled=True, loss_scale="dynamic", enable_caching=True, verbose=False, allow_banned=False):
+    global _DECORATOR_HANDLE
+
+    if not enabled:
+        handle = NoOpHandle()
+        _DECORATOR_HANDLE = handle
+        return handle
+
+    handle = AmpHandle(loss_scale, enable_caching, verbose)
+
+    # 0) Force-{fp16, fp32} for user-annotated functions
+    for mod, fn, cast_fn in _USER_CAST_REGISTRY:
+        try_caching = (cast_fn == utils.maybe_half)
+        wrap.cached_cast(mod, fn, cast_fn, handle,
+                         try_caching, verbose)
+    _USER_CAST_REGISTRY.clear()
+
+    # 0.5) Force-promote for user-annotated functions
+    for mod, fn in _USER_PROMOTE_REGISTRY:
+        wrap.promote(mod, fn, handle, verbose)
+    _USER_PROMOTE_REGISTRY.clear()
+
+    # 1) Force-{fp16, fp32} on white- / black-list functions
+    override_modules = [functional_overrides,
+                        torch_overrides,
+                        tensor_overrides]
+    cast_table = [('FP16_FUNCS', utils.maybe_half),
+                  ('FP32_FUNCS', utils.maybe_float)]
+    for module, (list_name, cast_fn) in itertools.product(override_modules,
+                                                          cast_table):
+        for fn in getattr(module, list_name):
+            try_caching = (cast_fn == utils.maybe_half)
+            wrap.cached_cast(module.MODULE, fn, cast_fn, handle,
+                             try_caching, verbose)
+
+    # 1.5) Pre-0.4, put the blacklist methods on HalfTensor and whitelist
+    #      methods on FloatTensor, since they're distinct types.
+    if compat.tensor_is_float_tensor():
+        for fn in tensor_overrides.FP16_FUNCS:
+            wrap.cached_cast(torch.cuda.FloatTensor, fn, utils.maybe_half,
+                             handle, try_caching=True, verbose=verbose)
+        for fn in tensor_overrides.FP32_FUNCS:
+            wrap.cached_cast(torch.cuda.HalfTensor, fn, utils.maybe_float,
+                             handle, try_caching=False, verbose=verbose)
+
+    # 2) Enable type-promotion on multi-arg functions and methods.
+    #    NB: special handling for sequence fns (e.g. `torch.cat`).
+    promote_modules = [torch_overrides, tensor_overrides]
+    promote_table = [('CASTS', wrap.promote),
+                     ('SEQUENCE_CASTS', wrap.sequence_promote)]
+    for promote_mod, (list_name, promote_fn) in itertools.product(promote_modules,
+                                                                  promote_table):
+        for fn in getattr(promote_mod, list_name):
+            promote_fn(promote_mod.MODULE, fn, handle, verbose)
+
+    # 2.5) Pre-0.4, add blacklist methods directly to HalfTensor and FloatTensor types
+    if compat.tensor_is_float_tensor():
+        for cls, (list_name, promote_fn) in itertools.product([torch.cuda.FloatTensor,
+                                                               torch.cuda.HalfTensor],
+                                                              promote_table):
+            for fn in getattr(tensor_overrides, list_name):
+                promote_fn(cls, fn, handle, verbose)
+
+    # 3) For any in-place version of a blacklist function, error if any input is fp16.
+    #    NB: this is overly conservative.
+    for fn in utils.as_inplace(torch_overrides.FP32_FUNCS):
+        wrap.err_if_any_half(torch_overrides.MODULE, fn, handle)
+
+    # 3.5) For any in-place blacklist method, error if called on fp16 tensor
+    for fn in utils.as_inplace(tensor_overrides.FP32_FUNCS):
+        wrap.err_if_arg0_half(tensor_overrides.MODULE, fn, handle, verbose)
+        if compat.tensor_is_float_tensor():
+            wrap.err_if_arg0_half(torch.cuda.HalfTensor, fn, handle, verbose)
+
+    # 4) For other in-place methods, match the type of self tensor
+    for fn in utils.as_inplace(itertools.chain(
+            tensor_overrides.FP16_FUNCS,
+            tensor_overrides.CASTS)):
+        wrap.promote_match_arg0(tensor_overrides.MODULE, fn, handle, verbose)
+        if compat.tensor_is_float_tensor():
+            wrap.promote_match_arg0(torch.cuda.HalfTensor, fn, handle, verbose)
+            wrap.promote_match_arg0(torch.cuda.FloatTensor, fn, handle, verbose)
+
+    # 5) RNNs + RNN cells are whitelisted specially
+    if rnn_compat.has_old_rnns():
+        wrap.rnn_cast(torch.nn.backends.thnn.backend, 'RNN', handle, verbose)
+    if not rnn_compat.has_old_rnns():
+        # Patch in our own indirection of `_VF` in modules/rnn s.t. it is mutable.
+        torch.nn.modules.rnn._VF = rnn_compat.VariableFunctionsShim()
+        # Wrap all the rnns
+        for x in rnn_compat.RNN_NAMES:
+            wrap.new_rnn_cast(x.upper(), handle, verbose)
+
+    # Wrap all the RNN cells
+    rnn_compat.whitelist_rnn_cells(handle, verbose)
+
+    # 6) Place error+print message on banned functions.
+    #    Or, if allow_banned, then cast to FP32.
+    for fn, err_msg in functional_overrides.BANNED_FUNCS:
+        if allow_banned:
+            wrap.cached_cast(functional_overrides.MODULE, fn, utils.maybe_float,
+                             handle, try_caching=True, verbose=verbose)
+        else:
+            wrap.err_if_any_half(functional_overrides.MODULE, fn, handle, err_msg)
+
+    _DECORATOR_HANDLE = handle
+
+    _amp_state.handle = handle
+
+    return handle
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/compat.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/compat.py
new file mode 100644
index 0000000000..22276bd47d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/compat.py
@@ -0,0 +1,46 @@
+import torch
+
+# True for post-0.4, when Variables/Tensors merged.
+def variable_is_tensor():
+    v = torch.autograd.Variable()
+    return isinstance(v, torch.Tensor)
+
+def tensor_is_variable():
+    x = torch.Tensor()
+    return type(x) == torch.autograd.Variable
+
+# False for post-0.4
+def tensor_is_float_tensor():
+    x = torch.Tensor()
+    return type(x) == torch.FloatTensor
+
+# Akin to `torch.is_tensor`, but returns True for Variable
+# objects in pre-0.4.
+def is_tensor_like(x):
+    return torch.is_tensor(x) or isinstance(x, torch.autograd.Variable)
+
+# Wraps `torch.is_floating_point` if present, otherwise checks
+# the suffix of `x.type()`.
+def is_floating_point(x):
+    if hasattr(torch, 'is_floating_point'):
+        return torch.is_floating_point(x)
+    try:
+        torch_type = x.type()
+        return torch_type.endswith('FloatTensor') or \
+            torch_type.endswith('HalfTensor') or \
+            torch_type.endswith('DoubleTensor')
+    except AttributeError:
+        return False
+
+def scalar_python_val(x):
+    if hasattr(x, 'item'):
+        return x.item()
+    else:
+        if isinstance(x, torch.autograd.Variable):
+            return x.data[0]
+        else:
+            return x[0]
+
+# Accounts for the possibility that some ops may be removed from a namespace.
+def filter_attrs(module, attrs):
+    return list(attrname for attrname in attrs if hasattr(module, attrname))
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/frontend.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/frontend.py
new file mode 100644
index 0000000000..da0f05dc99
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/frontend.py
@@ -0,0 +1,442 @@
+import torch
+from ._initialize import _initialize
+from ._amp_state import _amp_state, warn_or_err, maybe_print
+from collections import OrderedDict
+
+
+class Properties(object):
+    """
+    This class has two purposes: to establish a set of default properties,
+    and to route setting of these attributes through __setattr__ so that (in theory)
+    they can be checked for consistency with other existing args.
+    """
+    def __init__(self):
+        self.options = {
+            "enabled" : False,
+            "opt_level" : None,
+            "cast_model_type" : None,
+            "patch_torch_functions" : False,
+            "keep_batchnorm_fp32" : None,
+            "master_weights" : None,
+            "loss_scale" : 1.0,
+            # Reserved for future functionality
+            # "fused_optimizer" : False,
+            # "enable_ddp_interop" : False,
+            }
+
+    """
+    This function allows updating several options at a time without routing through
+    __setattr__ checks, to avoid "you can't get there from here" scenarios.
+    Currently not intended to be exposed; users are expected to select an opt_level
+    and apply consistent modifications.
+    """
+    def _update_options_dict(self, new_options):
+        for k, v in new_options:
+            if k in self.options:
+                self.options[k] = v
+            else:
+                raise ValueError("Tried to set unexpected option {}".format(k))
+    """
+    The members of "options" are not direct attributes of self, so access attempts
+    will roll down to __getattr__.  This borrows from the logic in torch.nn.Module.
+    """
+    def __getattr__(self, name):
+        if "options" in self.__dict__:
+            options =  self.__dict__["options"]
+            if name in options:
+                return options[name]
+        raise AttributeError("'{}' object has no attribute '{}'".format(
+            type(self).__name__, name))
+
+    def __setattr__(self, name, value):
+        if "options" in self.__dict__:
+            if name in self.options:
+                # print("setting {} {}".format(name, value))
+                if name == "cast_model_type":
+                    if self.opt_level == "O1" and value is not None:
+                        if value is not False:
+                            if value is not torch.float32:
+                                warn_or_err("O1 inserts casts around Torch functions rather than "
+                                            "model weights, so with O1, the model weights themselves "
+                                            "should remain FP32. If you wish to cast the model to a "
+                                            "different type, use opt_level='O2' or 'O3'. " +
+                                            "cast_model_type was {}".format(value))
+                    self.options[name] = value
+                elif name == "patch_torch_functions":
+                    if self.opt_level != "O1" and value:
+                        warn_or_err("Currently, patch_torch_functions=True should only be set by "
+                                    "selecting opt_level='O1'.")
+                    self.options[name] = value
+                elif name == "keep_batchnorm_fp32":
+                    if self.opt_level == "O1" and value is not None:
+                        warn_or_err("With opt_level O1, batchnorm functions are automatically patched "
+                                    "to run in FP32, so keep_batchnorm_fp32 should be None." +
+                                    " keep_batchnorm_fp32 was {}".format(value))
+                    if value == "False":
+                        self.options[name] = False
+                    elif value == "True":
+                        self.options[name] = True
+                    else:
+                        assert (value is True or value is False or value is None),\
+                            "keep_batchnorm_fp32 must be a boolean, the string 'True' or 'False', "\
+                            "or None, found keep_batchnorm_fp32={}".format(value)
+                        self.options[name] = value
+                elif name == "master_weights":
+                    if self.opt_level == "O1" and value is not None:
+                        warn_or_err("It doesn't make sense to use master_weights with O1. "
+                                    "With O1, your model weights themselves should be FP32.")
+                    self.options[name] = value
+                elif name == "loss_scale":
+                    if value == "dynamic":
+                        self.options[name] = value
+                    else:
+                        self.options[name] = float(value)
+                else:
+                    self.options[name] = value
+        else:
+            super(Properties, self).__setattr__(name, value)
+
+
+""" O0-O3 are convenience wrappers to establish defaults for typically used mixed precision options. """
+
+class O3:
+    brief = "O3:  Pure FP16 training."
+    more = "Calls .half() on your model, converting the entire model to FP16.\n"\
+        "A casting operation is also inserted to cast incoming Tensors to FP16,\n"\
+        "so you don't need to change your data pipeline.\n"\
+        "This mode is useful for establishing a performance ceiling.\n"\
+        "It's also possible training may 'just work' in this mode.\n"\
+        "If not, try other optimization levels."
+
+    def __call__(self, properties):
+        properties.enabled = True
+        properties.opt_level = "O3"
+        properties.cast_model_type = torch.float16
+        properties.patch_torch_functions = False
+        properties.keep_batchnorm_fp32 = False
+        properties.master_weights = False
+        properties.loss_scale = 1.0
+        # properties.fused_optimizer = False
+        # properties.enable_ddp_interop = False
+        return properties # modified in place so this isn't really necessary
+
+
+class O2:
+    brief = "O2:  FP16 training with FP32 batchnorm and FP32 master weights.\n"
+    more = "Calls .half() on your model, converting the entire model (except for batchnorms)\n"\
+        "to FP16.  Batchnorms are retained in FP32 for additional stability.\n"\
+        "The forward pass is patched to cast incoming Tensors to FP16, so you don't need to change\n"\
+        "your data pipeline.\n"\
+        "O2 creates FP32 master weights outside the model and patches any optimizers to update\n"\
+        "these master weights, then copy the master weights into the FP16 model weights.\n"\
+        "Master weights can also improve convergence and stability."
+
+    def __call__(self, properties):
+        properties.enabled = True
+        properties.opt_level = "O2"
+        properties.cast_model_type = torch.float16
+        properties.patch_torch_functions = False
+        properties.keep_batchnorm_fp32 = True
+        properties.master_weights = True
+        properties.loss_scale = "dynamic"
+        # properties.fused_optimizer = False
+        # properties.enable_ddp_interop = False
+        return properties # modified in place so this isn't really necessary
+
+
+class O1:
+    brief = "O1:  Insert automatic casts around Pytorch functions and Tensor methods.\n"
+    more = "The type of your model's weights is not altered.  However, internally,\n"\
+        "Pytorch functions are patched to cast any Tensor Core-friendly ops to FP16 for speed,\n"\
+        "while operations that might benefit from the additional stability of FP32 are patched\n"\
+        "to cast their inputs to fp32.\n"\
+        "O1 is the safest way to try mixed precision training, and is recommended when\n"\
+        "trying mixed precision training for the first time."
+
+    def __call__(self, properties):
+        properties.enabled = True
+        properties.opt_level = "O1"
+        properties.cast_model_type = None
+        properties.patch_torch_functions = True
+        properties.keep_batchnorm_fp32 = None
+        properties.master_weights = None
+        properties.loss_scale = "dynamic"
+        # properties.fused_optimizer = False
+        # properties.enable_ddp_interop = False
+        return properties # modified in place so this isn't really necessary
+
+
+class O0:
+    brief = "O0:  Pure FP32 training.\n"
+    more = "Your models are checked to make sure parameters are FP32, but otherwise the\n"\
+        "types of weights and internal Pytorch operations are not altered.  This mode disables any\n"\
+        "FP16 arithmetic, although other optimizations like DDP interop may still be requested.\n"
+
+    def __call__(self, properties):
+        properties.enabled = True
+        properties.opt_level = "O0"
+        properties.cast_model_type = torch.float32
+        properties.patch_torch_functions = False
+        properties.keep_batchnorm_fp32 = None
+        properties.master_weights = False
+        properties.loss_scale = 1.0
+        # properties.fused_optimizer = False
+        # properties.enable_ddp_interop = False
+        return properties # modified in place so this isn't really necessary
+
+
+opt_levels = {"O3": O3(),
+              "O2": O2(),
+              "O1": O1(),
+              "O0": O0()}
+
+
+# allow user to directly pass Properties struct as well?
+def initialize(
+    models,
+    optimizers=None,
+    enabled=True,
+    opt_level="O1",
+    cast_model_type=None,
+    patch_torch_functions=None,
+    keep_batchnorm_fp32=None,
+    master_weights=None,
+    loss_scale=None,
+    cast_model_outputs=None,
+    num_losses=1,
+    verbosity=1,
+    min_loss_scale=None,
+    max_loss_scale=2.**24
+    ):
+    """
+    Initialize your models, optimizers, and the Torch tensor and functional namespace according to the
+    chosen ``opt_level`` and overridden properties, if any.
+
+    ``amp.initialize`` should be called **after** you have finished
+    constructing your model(s) and
+    optimizer(s), but **before** you send your model through any DistributedDataParallel wrapper.
+    See `Distributed training`_ in the Imagenet example.
+
+    Currently, ``amp.initialize`` should only be called **once**,
+    although it can process an arbitrary number of
+    models and optimizers (see the corresponding `Advanced Amp Usage topic`_).
+    If you think your use case requires ``amp.initialize`` to be called more than once,
+    `let us know`_.
+
+    Any property keyword argument that is not ``None`` will be interpreted as a manual override.
+
+    To prevent having to rewrite anything else in your script, name the returned models/optimizers
+    to replace the passed models/optimizers, as in the code sample below.
+
+    Args:
+        models (torch.nn.Module or list of torch.nn.Modules):  Models to modify/cast.
+        optimizers (optional, torch.optim.Optimizer or list of torch.optim.Optimizers):  Optimizers to modify/cast.
+            REQUIRED for training, optional for inference.
+        enabled (bool, optional, default=True):  If False, renders all Amp calls no-ops, so your script
+            should run as if Amp were not present.
+        opt_level (str, optional, default="O1"):  Pure or mixed precision optimization level.  Accepted values are
+            "O0", "O1", "O2", and "O3", explained in detail above.
+        cast_model_type (``torch.dtype``, optional, default=None):  Optional property override, see
+            above.
+        patch_torch_functions (bool, optional, default=None):  Optional property override.
+        keep_batchnorm_fp32 (bool or str, optional, default=None):  Optional property override.  If
+            passed as a string, must be the string "True" or "False".
+        master_weights (bool, optional, default=None):  Optional property override.
+        loss_scale (float or str, optional, default=None):  Optional property override.  If passed as a string,
+            must be a string representing a number, e.g., "128.0", or the string "dynamic".
+        cast_model_outputs (torch.dtype, optional, default=None):  Option to ensure that the outputs
+            of your model(s) are always cast to a particular type regardless of ``opt_level``.
+        num_losses (int, optional, default=1):  Option to tell Amp in advance how many losses/backward
+            passes you plan to use.  When used in conjunction with the ``loss_id`` argument to
+            ``amp.scale_loss``, enables Amp to use a different loss scale per loss/backward pass,
+            which can improve stability.  See "Multiple models/optimizers/losses"
+            under `Advanced Amp Usage`_ for examples.  If ``num_losses`` is left to 1, Amp will still
+            support multiple losses/backward passes, but use a single global loss scale
+            for all of them.
+        verbosity (int, default=1):  Set to 0 to suppress Amp-related output.
+        min_loss_scale (float, default=None):  Sets a floor for the loss scale values that can be chosen by dynamic
+            loss scaling.  The default value of None means that no floor is imposed.
+            If dynamic loss scaling is not used, `min_loss_scale` is ignored.
+        max_loss_scale (float, default=2.**24):  Sets a ceiling for the loss scale values that can be chosen by
+            dynamic loss scaling.  If dynamic loss scaling is not used, `max_loss_scale` is ignored.
+
+    Returns:
+        Model(s) and optimizer(s) modified according to the ``opt_level``.
+        If either the ``models`` or ``optimizers`` args were lists, the corresponding return value will
+        also be a list.
+
+    Permissible invocations::
+
+        model, optim = amp.initialize(model, optim,...)
+        model, [optim1, optim2] = amp.initialize(model, [optim1, optim2],...)
+        [model1, model2], optim = amp.initialize([model1, model2], optim,...)
+        [model1, model2], [optim1, optim2] = amp.initialize([model1, model2], [optim1, optim2],...)
+
+        # This is not an exhaustive list of the cross product of options that are possible,
+        # just a set of examples.
+        model, optim = amp.initialize(model, optim, opt_level="O0")
+        model, optim = amp.initialize(model, optim, opt_level="O0", loss_scale="dynamic"|128.0|"128.0")
+
+        model, optim = amp.initialize(model, optim, opt_level="O1") # uses "loss_scale="dynamic" default
+        model, optim = amp.initialize(model, optim, opt_level="O1", loss_scale=128.0|"128.0")
+
+        model, optim = amp.initialize(model, optim, opt_level="O2") # uses "loss_scale="dynamic" default
+        model, optim = amp.initialize(model, optim, opt_level="O2", loss_scale=128.0|"128.0")
+        model, optim = amp.initialize(model, optim, opt_level="O2", keep_batchnorm_fp32=True|False|"True"|"False")
+
+        model, optim = amp.initialize(model, optim, opt_level="O3") # uses loss_scale=1.0 default
+        model, optim = amp.initialize(model, optim, opt_level="O3", loss_scale="dynamic"|128.0|"128.0")
+        model, optim = amp.initialize(model, optim, opt_level="O3", keep_batchnorm_fp32=True|False|"True"|"False")
+
+    The `Imagenet example`_ demonstrates live use of various opt_levels and overrides.
+
+    .. _`Distributed training`:
+        https://github.com/NVIDIA/apex/tree/master/examples/imagenet#distributed-training
+
+    .. _`Imagenet example`:
+        https://github.com/NVIDIA/apex/tree/master/examples/imagenet
+
+    .. _`Advanced Amp Usage`:
+        https://nvidia.github.io/apex/advanced.html
+
+    .. _`Advanced Amp Usage topic`:
+        https://nvidia.github.io/apex/advanced.html#multiple-models-optimizers-losses
+
+    .. _`let us know`:
+        https://github.com/NVIDIA/apex/issues
+    """
+    _amp_state.opt_properties = Properties()
+    _amp_state.verbosity = verbosity
+
+    if not enabled:
+        if optimizers is None:
+            return models
+        else:
+            return models, optimizers
+
+    if not torch.backends.cudnn.enabled:
+        raise RuntimeError(
+            "Amp requires torch.backends.cudnn.enabled = True")
+
+    if opt_level not in opt_levels:
+        raise RuntimeError(
+            "Unexpected optimization level {}. ".format(opt_level) +
+            "Options are 'O0', 'O1', 'O2', 'O3'.  Note that in `O0`, `O1`, etc., the prefix O is the letter O, " +
+            "not the number zero.")
+    else:
+        _amp_state.opt_properties = opt_levels[opt_level](_amp_state.opt_properties)
+        maybe_print("Selected optimization level {}".format(opt_levels[opt_level].brief), True)
+        maybe_print("Defaults for this optimization level are:", True)
+        for k, v in _amp_state.opt_properties.options.items():
+            maybe_print("{:22} : {}".format(k, v), True)
+
+    _amp_state.min_loss_scale = min_loss_scale
+    _amp_state.max_loss_scale = max_loss_scale
+
+    maybe_print("Processing user overrides (additional kwargs that are not None)...", True)
+    # I chose to have the keyword arguments listed directly in the argument list,
+    # instead of **kwargs, so I can't use kwargs.items() here.
+    if enabled is not None:
+        _amp_state.opt_properties.enabled = enabled
+    if opt_level is not None:
+        _amp_state.opt_properties.opt_level = opt_level
+    if cast_model_type is not None:
+        _amp_state.opt_properties.cast_model_type = cast_model_type
+    if patch_torch_functions is not None:
+        _amp_state.opt_properties.patch_torch_functions = patch_torch_functions
+    if keep_batchnorm_fp32 is not None:
+        _amp_state.opt_properties.keep_batchnorm_fp32 = keep_batchnorm_fp32
+    if master_weights is not None:
+        _amp_state.opt_properties.master_weights = master_weights
+    if loss_scale is not None:
+        _amp_state.opt_properties.loss_scale = loss_scale
+
+    maybe_print("After processing overrides, optimization options are:", True)
+    for k, v in _amp_state.opt_properties.options.items():
+        maybe_print("{:22} : {}".format(k, v), True)
+
+    return _initialize(models, optimizers, _amp_state.opt_properties, num_losses, cast_model_outputs)
+
+
+def state_dict(destination=None):
+    if destination is None:
+        destination = OrderedDict()
+
+    for idx, loss_scaler in enumerate(_amp_state.loss_scalers):
+        destination['loss_scaler%d' % idx] = {
+            'loss_scale': loss_scaler.loss_scale(),
+            'unskipped': loss_scaler._unskipped,
+        }
+    return destination
+
+
+def load_state_dict(state_dict):
+    # Check if state_dict containes the same number of loss_scalers as current setup
+    if len(state_dict) != len(_amp_state.loss_scalers):
+        print('Warning: state_dict contains {} entries, while {} loss_scalers are used'.format(
+            len(state_dict), len(_amp_state.loss_scalers)))
+
+    state_dict = state_dict.copy()
+    
+    nb_loss_scalers = len(_amp_state.loss_scalers)
+    unexpected_keys = []
+    # Initialize idx outside, since unexpected_keys will increase it if enumerate is used
+    idx = 0
+    for key in state_dict:
+        if 'loss_scaler' not in key:
+            unexpected_keys.append(key)
+        else:
+            if idx > (nb_loss_scalers - 1):
+                print('Skipping loss_scaler[{}], since num_losses was set to {}'.format(
+                    idx, nb_loss_scalers))
+                break
+            _amp_state.loss_scalers[idx]._loss_scale = state_dict[key]['loss_scale']
+            _amp_state.loss_scalers[idx]._unskipped = state_dict[key]['unskipped']
+            idx += 1
+
+    if len(unexpected_keys) > 0:
+        raise RuntimeError(
+            'Error(s) in loading state_dict. Unexpected key(s) in state_dict: {}. '.format(
+                ', '.join('"{}"'.format(k) for k in unexpected_keys)))
+
+
+# TODO:  is this necessary/useful?
+# def check_option_consistency(enabled=True,
+#                              opt_level=None,
+#                              cast_model_type=None,
+#                              patch_torch_functions=None,
+#                              keep_batchnorm_fp32=None,
+#                              master_weights=None,
+#                              loss_scale=None,
+#                              enable_ddp_interop=None,
+#                              hard_override=False):
+#     """
+#     Utility function that enables users to quickly check if the option combination they intend
+#     to use is permitted.  ``check_option_consistency`` does not require models or optimizers
+#     to be constructed, and can be called at any point in the script.  ``check_option_consistency``
+#     is totally self-contained; it does not set any amp global state or affect anything outside
+#     of itself.
+#     """
+#
+#     if not enabled:
+#         return
+#
+#     if opt_level not in opt_levels:
+#         raise RuntimeError("Unexpected optimization level.  Options are 'O0', 'O1', 'O2', 'O3'.")
+#     else:
+#         opt_properties = opt_levels[opt_level](Properties())
+#         print("Selected optimization level {}", opt_levels[opt_level].brief)
+#         print("Defaults for this optimization level are:")
+#         for k, v in opt_properties.options:
+#             print("{:22} : {}".format(k, v))
+#
+#     print("Processing user overrides (additional kwargs that are not None)...")
+#     for k, v in kwargs:
+#         if k not in _amp_state.opt_properties.options:
+#             raise RuntimeError("Unexpected kwarg {}".format(k))
+#         if v is not None:
+#             setattr(opt_properties, k, v)
+#
+#     print("After processing overrides, optimization options are:")
+#     for k, v in opt_properties.options:
+#         print("{:22} : {}".format(k, v))
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/handle.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/handle.py
new file mode 100644
index 0000000000..0be567ca48
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/handle.py
@@ -0,0 +1,281 @@
+import contextlib
+import warnings
+import sys
+import torch
+
+from . import utils
+from .opt import OptimWrapper
+from .scaler import LossScaler
+from ._amp_state import _amp_state, master_params, maybe_print
+
+if torch.distributed.is_available():
+    from ..parallel.LARC import LARC
+
+
+# There's no reason to expose the notion of a "handle". Everything can happen through amp.* calls.
+@contextlib.contextmanager
+def scale_loss(loss,
+               optimizers,
+               loss_id=0,
+               model=None,
+               delay_unscale=False,
+               delay_overflow_check=False):
+    """
+    On context manager entrance, creates ``scaled_loss = (loss.float())*current loss scale``.
+    ``scaled_loss`` is yielded so that the user can call ``scaled_loss.backward()``::
+
+        with amp.scale_loss(loss, optimizer) as scaled_loss:
+            scaled_loss.backward()
+
+    On context manager exit (if ``delay_unscale=False``), the gradients are checked for infs/NaNs
+    and unscaled, so that ``optimizer.step()`` can be called.
+
+    .. note::
+        If Amp is using explicit FP32 master params (which is the default for ``opt_level=O2``, and
+        can also be manually enabled by supplying ``master_weights=True`` to ``amp.initialize``)
+        any FP16 gradients are copied to FP32 master gradients before being unscaled.
+        ``optimizer.step()`` will then apply the unscaled master gradients to the master params.
+
+    .. warning::
+        If Amp is using explicit FP32 master params, only the FP32 master gradients will be
+        unscaled.  The direct ``.grad`` attributes of any FP16
+        model params will remain scaled after context manager exit.
+        This subtlety affects gradient clipping.  See "Gradient clipping" under
+        `Advanced Amp Usage`_ for best practices.
+
+    Args:
+        loss(Tensor):  Typically a scalar Tensor. The ``scaled_loss`` that the context
+            manager yields is simply ``loss.float()*loss_scale``, so in principle
+            ``loss`` could have more than one element, as long as you call
+            ``backward()`` on ``scaled_loss`` appropriately within the context manager body.
+        optimizers:  All optimizer(s) for which the current backward pass is creating gradients.
+            Must be an optimizer or list of optimizers returned from an earlier call
+            to ``amp.initialize``.  For example use with multiple optimizers, see
+            "Multiple models/optimizers/losses" under `Advanced Amp Usage`_.
+        loss_id(int, optional, default=0):  When used in conjunction with the ``num_losses`` argument
+            to ``amp.initialize``, enables Amp to use a different loss scale per loss.  ``loss_id``
+            must be an integer between 0 and ``num_losses`` that tells Amp which loss is
+            being used for the current backward pass.  See "Multiple models/optimizers/losses"
+            under `Advanced Amp Usage`_ for examples.  If ``loss_id`` is left unspecified, Amp
+            will use the default global loss scaler for this backward pass.
+        model(torch.nn.Module, optional, default=None):  Currently unused, reserved to enable future
+            optimizations.
+        delay_unscale(bool, optional, default=False):  ``delay_unscale`` is never necessary, and
+            the default value of ``False`` is strongly recommended.
+            If ``True``, Amp will not unscale the gradients or perform model->master
+            gradient copies on context manager exit.
+            ``delay_unscale=True`` is a minor ninja performance optimization and can result
+            in weird gotchas (especially with multiple models/optimizers/losses),
+            so only use it if you know what you're doing.
+            "Gradient accumulation across iterations" under `Advanced Amp Usage`_
+            illustrates a situation where this CAN (but does not need to) be used.
+
+    .. warning::
+        If ``delay_unscale`` is ``True`` for a given backward pass, ``optimizer.step()`` cannot be
+        called yet after context manager exit, and must wait for another, later backward context
+        manager invocation with ``delay_unscale`` left to False.
+
+    .. _`Advanced Amp Usage`:
+        https://nvidia.github.io/apex/advanced.html
+    """
+    if not hasattr(_amp_state, "opt_properties"):
+        raise RuntimeError("Invoked 'with amp.scale_loss`, but internal Amp state has not been initialized.  "
+                           "model, optimizer = amp.initialize(model, optimizer, opt_level=...) must be called "
+                           "before `with amp.scale_loss`.")
+
+    if not _amp_state.opt_properties.enabled:
+        yield loss
+        return
+
+    if isinstance(optimizers, torch.optim.Optimizer) or ('LARC' in globals() and isinstance(optimizers, LARC)):
+        optimizers = [optimizers]
+
+    loss_scaler = _amp_state.loss_scalers[loss_id]
+    loss_scale = loss_scaler.loss_scale()
+
+    if ((not _amp_state.opt_properties.master_weights)
+        and (not loss_scaler.dynamic)
+        and loss_scale == 1.0):
+        yield loss.float()
+        # Needing to drop the cache here as well is an ugly gotcha.
+        # But for now I think it's necessary to short-circuit.
+        # Probably ok to skip this if not delay_unscale
+        if _amp_state.opt_properties.patch_torch_functions:
+            _amp_state.handle._clear_cache()
+        return
+
+    if not delay_unscale:
+        if isinstance(optimizers, list):
+            for optimizer in optimizers:
+                if not optimizer._amp_stash.params_have_scaled_gradients:
+                    optimizer._prepare_amp_backward()
+
+    yield (loss.float())*loss_scale
+
+    if delay_unscale:
+        for optimizer in optimizers:
+            optimizer._amp_stash.params_have_scaled_gradients = True
+    else:
+        # FusedSGD may take care of unscaling as part of their step() methods.
+        # if not isinstance(optimizers, FP16_Optimizer_for_fused):
+            loss_scaler.clear_overflow_state()
+            for optimizer in optimizers:
+                optimizer._post_amp_backward(loss_scaler)
+                optimizer._amp_stash.params_have_scaled_gradients = False
+            # For future fused optimizers that enable sync-free dynamic loss scaling,
+            # should_skip will always be False.
+            should_skip = False if delay_overflow_check else loss_scaler.update_scale()
+            if should_skip:
+                for optimizer in optimizers:
+                    if not optimizer._amp_stash.already_patched:
+                        # Close on loss_scaler and loss_id as well, to be safe.  Probably not
+                        # necessary because amp.scale_loss is already creating a temporary scope.
+                        def patch_step(opt, loss_scaler, loss_id):
+                            opt_step = opt.step
+                            def skip_step(closure=None):
+                                if closure is not None:
+                                    raise RuntimeError("Currently, Amp does not support closure use with optimizers.")
+                                maybe_print(("Gradient overflow.  Skipping step, loss scaler " +
+                                             "{} reducing loss scale to {}").format(loss_id,
+                                             loss_scaler.loss_scale()))
+                                # TODO:  I don't like the special casing for different optimizer implementations.
+                                # Maybe skip should delegate to a method owned by the optimizers themselves.
+                                if hasattr(opt._amp_stash, "all_fp32_from_fp16_params"):
+                                    # Clear the master grads that wouldn't be zeroed by model.zero_grad()
+                                    for param in opt._amp_stash.all_fp32_from_fp16_params:
+                                        param.grad = None
+                                if hasattr(opt, "most_recent_scale"):
+                                    opt.most_recent_scale = 1.0
+                                    opt.scale_set_by_backward = False
+                                opt.step = opt_step
+                                opt._amp_stash.already_patched = False
+                            return skip_step
+                        optimizer.step = patch_step(optimizer, loss_scaler, loss_id)
+                        optimizer._amp_stash.already_patched = True
+
+    # Probably ok to skip this if not delay_unscale
+    if _amp_state.opt_properties.patch_torch_functions:
+        _amp_state.handle._clear_cache()
+
+
+# Free function version of AmpHandle.disable_casts, another step on the
+# path to removing the concept of "AmpHandle"
+@contextlib.contextmanager
+def disable_casts():
+    _amp_state.handle._is_active = False
+    yield
+    _amp_state.handle._is_active = True
+
+
+class AmpHandle(object):
+    def __init__(self, loss_scale="dynamic", enable_caching=True, verbose=False):
+        self._enable_caching = enable_caching
+        self._verbose = verbose
+        self._cache = dict()
+        self._default_scaler = LossScaler(loss_scale)
+        self._is_active = True
+        self._all_wrappers = []
+
+    def is_active(self):
+        return self._is_active
+
+    @contextlib.contextmanager
+    def _disable_casts(self):
+        self._is_active = False
+        yield
+        self._is_active = True
+
+    def wrap_optimizer(self, optimizer, num_loss=1):
+        self._default_scaler = None
+        return OptimWrapper(optimizer, self, num_loss)
+
+    @contextlib.contextmanager
+    def scale_loss(self, loss, optimizer):
+        raise RuntimeError("The old Amp API is no longer supported.  Please move to the new API, "
+            "documented here:  https://nvidia.github.io/apex/amp.html.  Transition guide:  "
+            "https://nvidia.github.io/apex/amp.html#transition-guide-for-old-api-users")
+
+        if not self.is_active():
+            yield loss
+            return
+
+        if self._default_scaler is None:
+            raise RuntimeError(
+                'After calling `handle.wrap_optimizer()`, you must explicitly ' +
+                'use `optimizer.scale_loss(loss)`.')
+
+        # TODO: this code block is duplicated here and `opt.py`. Unify.
+        loss_scale = self._default_scaler.loss_scale()
+        yield loss * loss_scale
+
+        self._default_scaler.clear_overflow_state()
+        self._default_scaler.unscale(
+            master_params(optimizer),
+            master_params(optimizer),
+            loss_scale)
+        should_skip = self._default_scaler.update_scale()
+        if should_skip:
+            optimizer_step = optimizer.step
+            def skip_step():
+                maybe_print('Gradient overflow, skipping update')
+                optimizer.step = optimizer_step
+            optimizer.step = skip_step
+
+        self._clear_cache()
+
+    def _clear_cache(self):
+        self._cache.clear()
+
+    # Experimental support for saving / restoring uncasted versions of functions
+    def _save_func(self, mod, fn, func):
+        self._all_wrappers.append((mod, fn, func))
+
+    def _deactivate(self):
+        for mod, fn, func in self._all_wrappers:
+            utils.set_func(mod, fn, func)
+        self._all_wrappers = []
+
+    @property
+    def has_cache(self):
+        return self._enable_caching
+
+    @property
+    def cache(self):
+        return self._cache
+
+    def remove_cache(self, param):
+        if self.has_cache and param in self.cache:
+            del self.cache[param]
+
+    @property
+    def verbose(self):
+        return self._verbose
+
+class NoOpHandle(object):
+    def is_active(self):
+        return False
+
+    @contextlib.contextmanager
+    def _disable_casts(self):
+        yield
+
+    def wrap_optimizer(self, optimizer, num_loss=1):
+        return OptimWrapper(optimizer, self, num_loss)
+
+    @contextlib.contextmanager
+    def scale_loss(self, loss, optimizer):
+        yield loss
+
+    @property
+    def has_cache(self):
+        return False
+
+    @property
+    def verbose(self):
+        return False
+
+    def _clear_cache(self):
+        pass
+
+    def _deactivate(self):
+        pass
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/functional_overrides.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/functional_overrides.py
new file mode 100644
index 0000000000..dd009cec6e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/functional_overrides.py
@@ -0,0 +1,80 @@
+
+# TODO: think about the following two. They do weird things.
+# - torch.nn.utils.clip_grad (but it should always be fp32 anyway)
+# - torch.nn.utils.weight_norm
+
+# Notes:
+# F.instance_norm uses batch_norm internally. Which correctly handles
+#   fp16 in/out with fp32 weights. So we shouldn't do anything for
+#   either of these.
+# F.normalize calls `input.norm()` internally, so it's redundant, but
+#   kept here in case impl. changes.
+# F.cosine_similarity is same: calls `x.norm()` internally.
+
+import torch.nn.functional
+
+MODULE = torch.nn.functional
+
+FP16_FUNCS = [
+    'conv1d',
+    'conv2d',
+    'conv3d',
+    'conv_transpose1d',
+    'conv_transpose2d',
+    'conv_transpose3d',
+    'conv_tbc', # Undocumented / maybe new?
+    'linear',
+]
+
+FP32_FUNCS = [
+
+    # Interpolation/Upsampling TODO:  Remove for 1.2
+    'interpolate',
+    'grid_sample',
+
+    # Pointwise
+    'softplus',
+    'softmin',
+    'log_softmax',
+    'softmax',
+    'gelu',
+    
+    # Normalization
+    'layer_norm',
+    'group_norm',
+    'local_response_norm',
+    'normalize',
+    'cosine_similarity',
+
+    # Loss functions
+    # TODO: which of these can be fp16?
+    'poisson_nll_loss',
+    'cosine_embedding_loss',
+    'cross_entropy',
+    'hinge_embedding_loss',
+    'kl_div',
+    'l1_loss',
+    'mse_loss',
+    'margin_ranking_loss',
+    'multilabel_margin_loss',
+    'multilabel_soft_margin_loss',
+    'multi_margin_loss',
+    'nll_loss',
+    'binary_cross_entropy_with_logits',
+    'smooth_l1_loss',
+    'soft_margin_loss',
+    'triplet_margin_loss',
+    'ctc_loss'
+]
+
+BANNED_FUNCS = [
+    ('binary_cross_entropy',
+     ("\namp does not work out-of-the-box with `F.binary_cross_entropy` or `torch.nn.BCELoss.` "
+      "It requires that the output of the previous function be already a FloatTensor. \n\n"
+      "Most models have a Sigmoid right before BCELoss. In that case, you can use\n"
+      "    torch.nn.BCEWithLogitsLoss\nto combine Sigmoid+BCELoss into a single layer "
+      "that is compatible with amp.\nAnother option is to add\n"
+      "    amp.register_float_function(torch, 'sigmoid')\nbefore calling `amp.init()`.\n"
+      "If you _really_ know what you are doing, you can disable this warning by passing "
+      "allow_banned=True to `amp.init()`."))
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/tensor_overrides.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/tensor_overrides.py
new file mode 100644
index 0000000000..18f3e5dcf2
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/tensor_overrides.py
@@ -0,0 +1,63 @@
+from .. import compat
+from . import torch_overrides
+
+import importlib
+
+import torch
+
+# if compat.variable_is_tensor() and not compat.tensor_is_variable():
+MODULE = torch.Tensor
+# else:
+#     MODULE = torch.autograd.Variable
+
+
+FP16_FUNCS = compat.filter_attrs(MODULE, [
+    '__matmul__',
+])
+
+FP32_FUNCS = compat.filter_attrs(MODULE, [
+    '__ipow__',
+    '__pow__',
+    '__rpow__',
+
+    # Cast to fp32 before transfer to CPU
+    'cpu',
+])
+
+CASTS = compat.filter_attrs(MODULE, [
+    '__add__',
+    '__div__',
+    '__eq__',
+    '__ge__',
+    '__gt__',
+    '__iadd__',
+    '__idiv__',
+    '__imul__',
+    '__isub__',
+    '__itruediv__',
+    '__le__',
+    '__lt__',
+    '__mul__',
+    '__ne__',
+    '__radd__',
+    '__rdiv__',
+    '__rmul__',
+    '__rsub__',
+    '__rtruediv__',
+    '__sub__',
+    '__truediv__',
+])
+
+# None of these, but here to make code cleaner.
+SEQUENCE_CASTS = []
+
+# We need to grab all the methods from torch_overrides and add them to
+# the Tensor lists as well, as almost all methods are duplicated
+# between `torch` and `torch.Tensor` (and check with `hasattr`,
+# because a few random ones aren't defined on Tensor)
+_self_mod = importlib.import_module(__name__)
+for attrname in ['FP16_FUNCS', 'FP32_FUNCS', 'CASTS', 'SEQUENCE_CASTS']:
+    lst = getattr(_self_mod, attrname)
+    for fn in getattr(torch_overrides, attrname):
+        if hasattr(MODULE, fn):
+            lst.append(fn)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/torch_overrides.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/torch_overrides.py
new file mode 100644
index 0000000000..7dedb05a83
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/lists/torch_overrides.py
@@ -0,0 +1,115 @@
+import torch
+
+from .. import utils
+
+MODULE = torch
+
+FP16_FUNCS = [
+    # Low level functions wrapped by torch.nn layers.
+    # The wrapper layers contain the weights which are then passed in as a parameter
+    # to these functions.
+    'conv1d',
+    'conv2d',
+    'conv3d',
+    'conv_transpose1d',
+    'conv_transpose2d',
+    'conv_transpose3d',
+    'conv_tbc',
+    'prelu',
+
+    # BLAS
+    'addmm',
+    'addmv',
+    'addr',
+    'matmul',
+    'mm',
+    'mv',
+]
+
+FP32_FUNCS = [
+    # Pointwise
+    'acos',
+    'asin',
+    'cosh',
+    'erfinv',
+    'exp',
+    'expm1',
+    'log',
+    'log10',
+    'log2',
+    'reciprocal',
+    'rsqrt',
+    'sinh',
+    'tan',
+
+    # Other math
+    'pow',
+
+    # Reduction
+    'cumprod',
+    'cumsum',
+    'dist',
+    # 'mean',
+    'norm',
+    'prod',
+    'std',
+    'sum',
+    'var',
+
+    # Misc
+    'renorm'
+]
+
+version_strings = torch.__version__.split('.')
+version_major = version_strings[0]
+version_minor = version_strings[1]
+version_num = float(version_major + "." + version_minor)
+# Before torch 1.1, mean must be blacklisted.
+if version_num < 1.1:
+    FP32_FUNCS.append('mean')
+
+# Before CUDA 9.1, batched matmul was missing fast FP16 kernels. We
+# check the CUDA version -- if at least 9.1, then put the bmm
+# functions on the fp16 list. Otherwise, put them on the fp32 list.
+_bmms = ['addbmm',
+         'baddbmm',
+         'bmm']
+
+if utils.is_cuda_enabled():
+  # workaround https://github.com/facebookresearch/maskrcnn-benchmark/issues/802
+  if utils.get_cuda_version() >= (9, 1, 0):
+      FP16_FUNCS.extend(_bmms)
+  else:
+      FP32_FUNCS.extend(_bmms)
+
+# Multi-tensor fns that may need type promotion
+CASTS = [
+    # Multi-tensor math
+    'addcdiv',
+    'addcmul',
+    'atan2',
+    'cross',
+    'bilinear',
+    'dot',
+
+    # Element-wise _or_ tensor-wise math
+    'add',
+    'div',
+    'mul',
+
+    # Comparison
+    'eq',
+    'equal',
+    'ge',
+    'gt',
+    'le',
+    'lt',
+    'ne'
+]
+
+# Functions that take sequence arguments. We need to inspect the whole
+# sequence and cast to the widest type.
+SEQUENCE_CASTS = [
+    'cat',
+    'stack'
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/opt.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/opt.py
new file mode 100644
index 0000000000..baf311684d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/opt.py
@@ -0,0 +1,103 @@
+import contextlib
+import warnings
+
+from .scaler import LossScaler, master_params
+from ._amp_state import maybe_print
+
+import numpy as np
+
+class OptimWrapper(object):
+    def __init__(self, optimizer, amp_handle, num_loss):
+        self._optimizer = optimizer
+        self._amp_handle = amp_handle
+        self._num_loss = num_loss
+        self._loss_idx = 0
+        self._skip_next = [False] * num_loss
+        self._loss_scaler = [LossScaler('dynamic') for _ in range(num_loss)]
+
+    @contextlib.contextmanager
+    def scale_loss(self, loss):
+        if not self._amp_handle.is_active():
+            yield loss
+            return
+
+        # When there are multiple losses per-optimizer, we need
+        # to save out current grad accumulation, since we won't be
+        # able to unscale this particulare loss once the grads are
+        # all mixed together.
+        cached_grads = []
+        if self._loss_idx > 0:
+            for p in master_params(self._optimizer):
+                if p.grad is not None:
+                    cached_grads.append(p.grad.data.detach().clone())
+                else:
+                    cached_grads.append(None)
+            self._optimizer.zero_grad()
+
+        loss_scale = self._cur_loss_scaler().loss_scale()
+        yield loss * loss_scale
+
+        self._cur_loss_scaler().clear_overflow_state()
+        self._cur_loss_scaler().unscale(
+            master_params(self._optimizer),
+            master_params(self._optimizer),
+            loss_scale)
+        self._skip_next[self._loss_idx] = self._cur_loss_scaler().update_scale()
+        self._loss_idx += 1
+
+        if len(cached_grads) > 0:
+            for p, cached_grad in zip(master_params(self._optimizer),
+                                      cached_grads):
+                if cached_grad is not None:
+                    p.grad.data.add_(cached_grad)
+            cached_grads = []
+
+    def _cur_loss_scaler(self):
+        assert 0 <= self._loss_idx < self._num_loss
+        return self._loss_scaler[self._loss_idx]
+
+    def step(self, closure=None):
+        if not self._amp_handle.is_active():
+            return self._optimizer.step(closure=closure)
+
+        self._loss_idx = 0
+
+        for group in self._optimizer.param_groups:
+            for p in group['params']:
+                self._amp_handle.remove_cache(p)
+
+        if closure is not None:
+            raise NotImplementedError(
+                'The `closure` argument is unsupported by the amp ' +
+                'optimizer wrapper.')
+        if any(self._skip_next):
+            maybe_print('Gradient overflow, skipping update')
+            self._skip_next = [False] * self._num_loss
+        else:
+            return self._optimizer.step(closure=closure)
+
+    # Forward any attribute lookups
+    def __getattr__(self, attr):
+        return getattr(self._optimizer, attr)
+
+    # Forward all torch.optim.Optimizer methods
+    def __getstate__(self):
+        return self._optimizer.__getstate__()
+
+    def __setstate__(self):
+        return self._optimizer.__setstate__()
+
+    def __repr__(self):
+        return self._optimizer.__repr__()
+
+    def state_dict(self):
+        return self._optimizer.state_dict()
+
+    def load_state_dict(self, state_dict):
+        return self._optimizer.load_state_dict(state_dict)
+
+    def zero_grad(self):
+        return self._optimizer.zero_grad()
+
+    def add_param_group(self, param_group):
+        return self._optimizer.add_param_group(param_group)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/rnn_compat.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/rnn_compat.py
new file mode 100644
index 0000000000..d062ae2658
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/rnn_compat.py
@@ -0,0 +1,53 @@
+from . import utils, wrap
+
+import torch
+_VF = torch._C._VariableFunctions
+RNN_NAMES = ['rnn_relu', 'rnn_tanh', 'gru', 'lstm']
+
+def _gen_VF_wrapper(name):
+    def wrapper(*args, **kwargs):
+        return getattr(_VF, name)(*args, **kwargs)
+    return wrapper
+
+# Some python magic to generate an object that has the rnn cell functions
+# defined on it, all of which call into corresponding _VF version.
+# Intended to patch torch.nn.modules.rnn._VF (aka, the ref named "_VF"
+# imported at module scope within torch.nn.modules.rnn).  This should
+# not affect third-party importers of _VF.py.
+class VariableFunctionsShim(object):
+    def __init__(self):
+        for name in RNN_NAMES:
+            for suffix in ['', '_cell']:
+               fn_name = name + suffix
+               setattr(self, fn_name, _gen_VF_wrapper(fn_name))
+
+def has_old_rnns():
+    try:
+        torch.nn.backends.thnn.backend.LSTMCell
+        return True
+    except:
+        return False
+
+def whitelist_rnn_cells(handle, verbose):
+    # Different module + function names in old/new RNN cases
+    if has_old_rnns():
+        fn_names = ['RNNReLUCell', 'RNNTanhCell', 'LSTMCell', 'GRUCell']
+        mod = torch.nn.backends.thnn.backend
+    else:
+        fn_names = [x + '_cell' for x in RNN_NAMES]
+        mod = torch.nn.modules.rnn._VF
+        assert isinstance(mod, VariableFunctionsShim)
+
+    # Insert casts on cell functions
+    for fn in fn_names:
+        wrap.cached_cast(mod, fn, utils.maybe_half, handle,
+                         try_caching=True, verbose=verbose)
+
+    if has_old_rnns():
+        # Special handling of `backward` for fused gru / lstm:
+        # The `backward` method calls Tensor.sum() (blacklist) internally,
+        # and then the resulting grad_input has the wrong type.
+        # TODO: where else is this a problem?
+        for rnn_type in ['GRUFused', 'LSTMFused']:
+            mod = getattr(torch.nn._functions.thnn.rnnFusedPointwise, rnn_type)
+            wrap.disable_casts(mod, 'backward', handle)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/scaler.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/scaler.py
new file mode 100644
index 0000000000..99888bc6fc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/scaler.py
@@ -0,0 +1,217 @@
+import torch
+from ..multi_tensor_apply import multi_tensor_applier
+from ._amp_state import _amp_state, master_params, maybe_print
+from itertools import product
+
+def scale_check_overflow_python(model_grad, master_grad, scale, check_overflow=False):
+    # Exception handling for 18.04 compatibility
+    if check_overflow:
+        cpu_sum = float(model_grad.float().sum())
+        if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum:
+            return True
+
+    if master_grad is not model_grad: # copy_ probably internally short-circuits this
+        master_grad.copy_(model_grad)
+    if scale != 1.0:
+        master_grad.mul_(scale)
+    return False
+
+def axpby_check_overflow_python(model_grad, stashed_grad, master_grad, a, b, check_overflow=False):
+    # Exception handling for 18.04 compatibility
+    if check_overflow:
+        cpu_sum = float(model_grad.float().sum())
+        if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum:
+            return True
+
+    # if master_grad is not model_grad: # copy_ probably internally short-circuits this
+    #     master_grad.copy_(model_grad)
+    assert stashed_grad.dtype == master_grad.dtype
+    converted_model_grad = model_grad.data.to(master_grad.dtype)
+    master_grad.data = a*converted_model_grad.data + b*stashed_grad.data
+    return False
+
+class LossScaler(object):
+    warned_no_fused_kernel = False
+    warned_unscaling_non_fp32_grad = False
+    has_fused_kernel = False
+
+    def __init__(self,
+                 loss_scale,
+                 init_scale=2.**16,
+                 scale_factor=2.,
+                 scale_window=2000,
+                 min_loss_scale=None,
+                 max_loss_scale=2.**24):
+        if loss_scale == "dynamic":
+            self.dynamic = True
+            self._loss_scale = min(max_loss_scale, init_scale)
+        else:
+            self.dynamic = False
+            self._loss_scale = loss_scale
+        self._max_loss_scale = max_loss_scale
+        self._min_loss_scale = min_loss_scale
+        self._scale_seq_len = scale_window
+        self._unskipped = 0
+        self._has_overflow = False
+        self._overflow_buf = torch.cuda.IntTensor([0])
+        if multi_tensor_applier.available:
+            import amp_C
+            LossScaler.has_fused_kernel = multi_tensor_applier.available
+            LossScaler.multi_tensor_scale_cuda = amp_C.multi_tensor_scale
+            LossScaler.multi_tensor_axpby_cuda = amp_C.multi_tensor_axpby
+        else:
+            if not LossScaler.warned_no_fused_kernel:
+                maybe_print(
+                    "Warning:  multi_tensor_applier fused unscale kernel is unavailable, "
+                    "possibly because apex was installed without --cuda_ext --cpp_ext. "
+                    "Using Python fallback.  Original ImportError was: " +
+                    repr(multi_tensor_applier.import_err),
+                    True)
+            LossScaler.has_fused_kernel = False
+            LossScaler.warned_no_fused_kernel = True
+
+    def loss_scale(self):
+        return self._loss_scale
+
+    def unscale_python(self, model_grads, master_grads, scale):
+        for model, master in zip(model_grads, master_grads):
+            if model is not None:
+                if not LossScaler.warned_unscaling_non_fp32_grad:
+                    if master.dtype != torch.float32:
+                        maybe_print(
+                            "Attempting to unscale a grad with type {} ".format(master.type()) +
+                            "Unscaling non-fp32 grads may indicate an error. "
+                            "When using Amp, you don't need to call .half() on your model.")
+                        LossScaler.warned_unscaling_non_fp32_grad = True
+                self._has_overflow = scale_check_overflow_python(model,
+                                                                 master,
+                                                                 1./scale,
+                                                                 self.dynamic)
+                if self._has_overflow and self.dynamic:
+                    break
+
+    # unused_scale keeps some of the old API alive for hopefully a short time.
+    def unscale(self, model_grads, master_grads, unused_scale, models_are_masters=False, scale_override=None):
+        if self._has_overflow:
+            return
+
+        scale = self._loss_scale
+        if scale_override is not None:
+            scale = scale_override
+
+        if scale == 1.0 and models_are_masters and not self.dynamic:
+            return
+
+        if LossScaler.has_fused_kernel:
+            # if (not LossScaler.warned_unscaling_non_fp32_grad
+            #     and master_grads[0].dtype == torch.float16):
+            #     print("Warning:  unscaling grads that are not FP32. "
+            #           "Unscaling non-fp32 grads may indicate an error. "
+            #           "When using Amp, you don't need to call .half() on your model.")
+            #     # Setting this to True unconditionally allows the possibility of an escape
+            #     # if never-before-seen non-fp32 grads are created in some later iteration.
+            #     LossScaler.warned_unscaling_non_fp32_grad = True
+            multi_tensor_applier(LossScaler.multi_tensor_scale_cuda,
+                                 self._overflow_buf,
+                                 [model_grads, master_grads],
+                                 1./scale)
+        else:
+            self.unscale_python(model_grads, master_grads, scale)
+
+        # Defer to update_scale
+        # If the fused kernel is available, we only need one D2H memcopy and sync.
+        # if LossScaler.has_fused_kernel and self.dynamic and not self._has_overflow:
+        #     self._has_overflow = self._overflow_buf.item()
+
+    def unscale_with_stashed_python(self,
+                                    model_grads,
+                                    stashed_master_grads,
+                                    master_grads,
+                                    a,
+                                    b):
+        for model, stashed, master in zip(model_grads, stashed_master_grads, master_grads):
+            if model is None and stashed is None:
+                continue
+            else:
+                if not LossScaler.warned_unscaling_non_fp32_grad:
+                    if master.dtype != torch.float32:
+                        maybe_print(
+                            "Attempting to unscale a grad with type {} ".format(master.type()) +
+                            "Unscaling non-fp32 grads may indicate an error. "
+                            "When using Amp, you don't need to call .half() on your model.")
+                        LossScaler.warned_unscaling_non_fp32_grad = True
+                self._has_overflow = axpby_check_overflow_python(model,
+                                                                 stashed,
+                                                                 master,
+                                                                 a,
+                                                                 b,
+                                                                 self.dynamic)
+                if self._has_overflow and self.dynamic:
+                    break
+
+    def unscale_with_stashed(self,
+                             model_grads,
+                             stashed_master_grads,
+                             master_grads,
+                             scale_override=None):
+        if self._has_overflow:
+            return
+
+        grads_have_scale, stashed_have_scale, out_scale = self._loss_scale, 1.0, 1.0
+        if scale_override is not None:
+            grads_have_scale, stashed_have_scale, out_scale = scale_override
+
+        if LossScaler.has_fused_kernel:
+            if (not LossScaler.warned_unscaling_non_fp32_grad
+                and master_grads[0].dtype == torch.float16):
+                print("Warning:  unscaling grads that are not FP32. "
+                      "Unscaling non-fp32 grads may indicate an error. "
+                      "When using Amp, you don't need to call .half() on your model.")
+                # Setting this to True unconditionally allows the possibility of an escape
+                # if never-before-seen non-fp32 grads are created in some later iteration.
+                LossScaler.warned_unscaling_non_fp32_grad = True
+            multi_tensor_applier(LossScaler.multi_tensor_axpby_cuda,
+                                 self._overflow_buf,
+                                 [model_grads, stashed_master_grads, master_grads],
+                                 out_scale/grads_have_scale,   # 1./scale,
+                                 out_scale/stashed_have_scale, # 1.0,
+                                 0) # check only arg 0, aka the incoming model grads, for infs
+        else:
+            self.unscale_with_stashed_python(model_grads,
+                                             stashed_master_grads,
+                                             master_grads,
+                                             out_scale/grads_have_scale,
+                                             out_scale/stashed_have_scale)
+
+        # Defer to update_scale
+        # If the fused kernel is available, we only need one D2H memcopy and sync.
+        # if LossScaler.has_fused_kernel and self.dynamic and not self._has_overflow:
+        #     self._has_overflow = self._overflow_buf.item()
+
+    def clear_overflow_state(self):
+        self._has_overflow = False
+        if self.has_fused_kernel:
+            self._overflow_buf.zero_()
+
+    # Separate so unscale() can be called more that once before updating.
+    def update_scale(self):
+        # If the fused kernel is available, we only need one D2H memcopy and sync.
+        if LossScaler.has_fused_kernel and self.dynamic and not self._has_overflow:
+            self._has_overflow = self._overflow_buf.item()
+
+        if self._has_overflow and self.dynamic:
+            should_skip = True
+            if(self._min_loss_scale):
+                self._loss_scale = max(self._min_loss_scale, self._loss_scale/2.)
+            else:
+                self._loss_scale = self._loss_scale/2.
+            self._unskipped = 0
+        else:
+            should_skip = False
+            self._unskipped += 1
+
+        if self._unskipped == self._scale_seq_len and self.dynamic:
+            self._loss_scale = min(self._max_loss_scale, self._loss_scale*2.)
+            self._unskipped = 0
+
+        return should_skip
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/utils.py
new file mode 100644
index 0000000000..0590cd70a1
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/utils.py
@@ -0,0 +1,210 @@
+from . import compat
+
+import functools
+import itertools
+
+import torch
+
+def is_cuda_enabled():
+    return torch.version.cuda is not None
+
+def get_cuda_version():
+    return tuple(int(x) for x in torch.version.cuda.split('.'))
+
+def is_fp_tensor(x):
+    if is_nested(x):
+        # Fast-fail version of all(is_fp_tensor)
+        for y in x:
+            if not is_fp_tensor(y):
+                return False
+        return True
+    return compat.is_tensor_like(x) and compat.is_floating_point(x)
+
+def is_nested(x):
+    return isinstance(x, tuple) or isinstance(x, list)
+
+def should_cache(x):
+    if is_nested(x):
+        # Fast-fail version of all(should_cache)
+        for y in x:
+            if not should_cache(y):
+                return False
+        return True
+    return isinstance(x, torch.nn.parameter.Parameter) and \
+        type_string(x) == 'FloatTensor'
+
+def collect_fp_tensor_types(args, kwargs):
+    def collect_types(x, types):
+        if is_nested(x):
+            for y in x:
+                collect_types(y, types)
+        else:
+            types.add(type_string(x))
+
+    all_args = itertools.chain(args, kwargs.values())
+    types = set()
+    for x in all_args:
+        if is_fp_tensor(x):
+            collect_types(x, types)
+    return types
+
+def type_string(x):
+    return x.type().split('.')[-1]
+
+def maybe_half(x, name='', verbose=False):
+    if is_nested(x):
+        return type(x)([maybe_half(y) for y in x])
+
+    if not x.is_cuda or type_string(x) == 'HalfTensor':
+        return x
+    else:
+        if verbose:
+            print('Float->Half ({})'.format(name))
+        return x.half()
+
+def maybe_float(x, name='', verbose=False):
+    if is_nested(x):
+        return type(x)([maybe_float(y) for y in x])
+
+    if not x.is_cuda or type_string(x) == 'FloatTensor':
+        return x
+    else:
+        if verbose:
+            print('Half->Float ({})'.format(name))
+        return x.float()
+
+# NB: returneds casted `args`, mutates `kwargs` in-place
+def casted_args(cast_fn, args, kwargs):
+    new_args = []
+    for x in args:
+        if is_fp_tensor(x):
+            new_args.append(cast_fn(x))
+        else:
+            new_args.append(x)
+    for k in kwargs:
+        val = kwargs[k]
+        if is_fp_tensor(val):
+            kwargs[k] = cast_fn(val)
+    return new_args
+
+def cached_cast(cast_fn, x, cache):
+    if is_nested(x):
+        return type(x)([cached_cast(y) for y in x])
+    if x in cache:
+        cached_x = cache[x]
+        if x.requires_grad and cached_x.requires_grad:
+            # Make sure x is actually cached_x's autograd parent.
+            if cached_x.grad_fn.next_functions[1][0].variable is not x:
+                raise RuntimeError("x and cache[x] both require grad, but x is not "
+                                   "cache[x]'s parent.  This is likely an error.")
+        # During eval, it's possible to end up caching casted weights with
+        # requires_grad=False.  On the next training iter, if cached_x is found
+        # and reused from the cache, it will not actually have x as its parent.
+        # Therefore, we choose to invalidate the cache (and force refreshing the cast)
+        # if x.requires_grad and cached_x.requires_grad do not match.
+        #
+        # During eval (i.e. running under with torch.no_grad()) the invalidation
+        # check would cause the cached value to be dropped every time, because
+        # cached_x would always be created with requires_grad=False, while x would
+        # still have requires_grad=True.  This would render the cache effectively
+        # useless during eval.  Therefore, if we are running under the no_grad()
+        # context manager (torch.is_grad_enabled=False) we elide the invalidation
+        # check, and use the cached value even though its requires_grad flag doesn't
+        # match.  During eval, we don't care that there's no autograd-graph
+        # connection between x and cached_x.
+        if torch.is_grad_enabled() and x.requires_grad != cached_x.requires_grad:
+            del cache[x]
+        else:
+            return cached_x
+
+    casted_x = cast_fn(x)
+    cache[x] = casted_x
+    return casted_x
+
+def verbosify(cast_fn, fn_name, verbose):
+    if verbose:
+        return functools.partial(cast_fn, name=fn_name, verbose=verbose)
+    else:
+        return cast_fn
+
+def as_inplace(fns):
+    for x in fns:
+        yield x + '_'
+
+def has_func(mod, fn):
+    if isinstance(mod, dict):
+        return fn in mod
+    else:
+        return hasattr(mod, fn)
+
+def get_func(mod, fn):
+    if isinstance(mod, dict):
+        return mod[fn]
+    else:
+        return getattr(mod, fn)
+
+def set_func(mod, fn, new_fn):
+    if isinstance(mod, dict):
+        mod[fn] = new_fn
+    else:
+        setattr(mod, fn, new_fn)
+
+def set_func_save(handle, mod, fn, new_fn):
+    cur_fn = get_func(mod, fn)
+    handle._save_func(mod, fn, cur_fn)
+    set_func(mod, fn, new_fn)
+
+# A couple problems get solved here:
+# - The flat_weight buffer is disconnected from autograd graph,
+#   so the fp16 weights need to be derived from the input weights
+#   to this forward call, not the flat buffer.
+# - The ordering of weights in the flat buffer is...idiosyncratic.
+# First problem is solved with combination of set_ (to set up
+# correct storage) and copy_ (so the fp16 weight derives from the
+# fp32 one in autograd.
+# Second is solved by doing ptr arithmetic on the fp32 weights
+# to derive the correct offset.
+#
+# TODO: maybe this should actually use
+# `torch._cudnn_rnn_flatten_weight`? But then I need to call
+# on first iter and cache the right offsets. Ugh.
+def synthesize_flattened_rnn_weights(fp32_weights,
+                                     fp16_flat_tensor,
+                                     rnn_fn='',
+                                     verbose=False):
+    fp16_weights = []
+    fp32_base_ptr = fp32_weights[0][0].data_ptr()
+    for layer_weights in fp32_weights:
+        fp16_layer_weights = []
+        for w_fp32 in layer_weights:
+            w_fp16 = w_fp32.new().half()
+            offset = (w_fp32.data_ptr() - fp32_base_ptr) // w_fp32.element_size()
+            w_fp16.set_(fp16_flat_tensor.storage(),
+                        offset,
+                        w_fp32.shape)
+            w_fp16.copy_(w_fp32)
+            if verbose:
+                print('Float->Half ({})'.format(rnn_fn))
+            fp16_layer_weights.append(w_fp16)
+        fp16_weights.append(fp16_layer_weights)
+    return fp16_weights
+
+# Roughly same as above, just the `fp32_weights` aren't nested.
+# Code kept separate for readability.
+def new_synthesize_flattened_rnn_weights(fp32_weights,
+                                         fp16_flat_tensor,
+                                         rnn_fn='',
+                                         verbose=False):
+    fp16_weights = []
+    fp32_base_ptr = fp32_weights[0].data_ptr()
+    for w_fp32 in fp32_weights:
+        w_fp16 = w_fp32.new().half()
+        offset = (w_fp32.data_ptr() - fp32_base_ptr) // w_fp32.element_size()
+        w_fp16.set_(fp16_flat_tensor.storage(),
+                    offset,
+                    w_fp32.shape)
+        w_fp16.copy_(w_fp32)
+        if verbose:
+            print('Float->Half ({})'.format(rnn_fn))
+        fp16_weights.append(w_fp16)
+    return fp16_weights
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/wrap.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/wrap.py
new file mode 100644
index 0000000000..559d0558d9
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/amp/wrap.py
@@ -0,0 +1,276 @@
+from . import compat
+from . import utils
+from ._amp_state import _amp_state
+from . import rnn_compat
+
+import functools
+
+import torch
+
+def make_cast_wrapper(orig_fn, cast_fn, handle,
+                      try_caching=False):
+    @functools.wraps(orig_fn)
+    def wrapper(*args, **kwargs):
+        if not handle.is_active():
+            return orig_fn(*args, **kwargs)
+
+        if try_caching and handle.has_cache:
+            args = list(args)
+            for i in range(len(args)):
+                if utils.should_cache(args[i]):
+                    args[i] = utils.cached_cast(cast_fn, args[i], handle.cache)
+            for k in kwargs:
+                if utils.should_cache(kwargs[k]):
+                    kwargs[k] = utils.cached_cast(cast_fn, kwargs[k], handle.cache)
+        new_args = utils.casted_args(cast_fn,
+                                     args,
+                                     kwargs)
+        return orig_fn(*new_args, **kwargs)
+    return wrapper
+
+def cached_cast(mod, fn, cast_fn, handle,
+                try_caching=False, verbose=False):
+    if not utils.has_func(mod, fn):
+        return
+
+    orig_fn = utils.get_func(mod, fn)
+    cast_fn = utils.verbosify(cast_fn, fn, verbose)
+    wrapper = make_cast_wrapper(orig_fn, cast_fn, handle, try_caching)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+# `handle` arg is unused, but simplifies API to make `make_cast_wrapper`
+# Annoyingly, make_promote_wrapper still uses the global handle.  Once everyone
+# is on the new API and I am free to get rid of handle, I can clean this up.
+def make_promote_wrapper(orig_fn, cast_fn, handle=None):
+    @functools.wraps(orig_fn)
+    def wrapper(*args, **kwargs):
+        if not _amp_state.handle.is_active():
+            return orig_fn(*args, **kwargs)
+
+        types = utils.collect_fp_tensor_types(args, kwargs)
+
+        if len(types) <= 1:
+            return orig_fn(*args, **kwargs)
+        elif len(types) == 2 and types == set(['HalfTensor', 'FloatTensor']):
+            new_args = utils.casted_args(cast_fn,
+                                         args,
+                                         kwargs)
+            return orig_fn(*new_args, **kwargs)
+        else:
+            raise NotImplementedError('Do not know how to handle ' +
+                                      'these types to promote: {}'
+                                      .format(types))
+    return wrapper
+
+def promote(mod, fn, handle, verbose=False):
+    orig_fn = utils.get_func(mod, fn)
+    maybe_float = utils.verbosify(utils.maybe_float, fn, verbose)
+    wrapper = make_promote_wrapper(orig_fn, maybe_float)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+def sequence_promote(mod, fn, handle, verbose=False):
+    orig_fn = utils.get_func(mod, fn)
+    maybe_float = utils.verbosify(utils.maybe_float, fn, verbose)
+    @functools.wraps(orig_fn)
+    def wrapper(seq, *args, **kwargs):
+        if not _amp_state.handle.is_active():
+            return orig_fn(seq, *args, **kwargs)
+
+        types = set([utils.type_string(x) for x in seq])
+        if len(types) <= 1:
+            return orig_fn(seq, *args, **kwargs)
+        elif types == set(['HalfTensor', 'FloatTensor']):
+            cast_seq = utils.casted_args(maybe_float,
+                                         seq, {})
+            return orig_fn(cast_seq, *args, **kwargs)
+        else:
+            # TODO: other mixed-type cases aren't due to amp.
+            #       Just pass through?
+            return orig_fn(seq, *args, **kwargs)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+def promote_match_arg0(mod, fn, handle, verbose=False):
+    if not utils.has_func(mod, fn):
+        return
+
+    orig_fn = utils.get_func(mod, fn)
+    @functools.wraps(orig_fn)
+    def wrapper(arg0, *args, **kwargs):
+        assert compat.is_tensor_like(arg0)
+        if not _amp_state.handle.is_active():
+            return orig_fn(arg0, *args, **kwargs)
+
+        if utils.type_string(arg0) == 'HalfTensor':
+            cast_fn = utils.maybe_half
+        elif utils.type_string(arg0) == 'FloatTensor':
+            cast_fn = utils.maybe_float
+        else:
+            return orig_fn(arg0, *args, **kwargs)
+        cast_fn = utils.verbosify(cast_fn, fn, verbose)
+        new_args = utils.casted_args(cast_fn, args, kwargs)
+        return orig_fn(arg0, *new_args, **kwargs)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+def err_if_any_half(mod, fn, handle, custom_err_msg=None):
+    if not utils.has_func(mod, fn):
+        return
+
+    orig_fn = utils.get_func(mod, fn)
+    @functools.wraps(orig_fn)
+    def wrapper(*args, **kwargs):
+        types = utils.collect_fp_tensor_types(args, kwargs)
+        if 'HalfTensor' in types:
+            if custom_err_msg:
+                raise NotImplementedError(custom_err_msg)
+            else:
+                raise NotImplementedError('Cannot call in-place function ' +
+                                          '{} with fp16 arguments.'.format(fn))
+        else:
+            return orig_fn(*args, **kwargs)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+def err_if_arg0_half(mod, fn, handle, verbose=False):
+    if not utils.has_func(mod, fn):
+        return
+
+    orig_fn = utils.get_func(mod, fn)
+    @functools.wraps(orig_fn)
+    def wrapper(arg0, *args, **kwargs):
+        assert compat.is_tensor_like(arg0)
+        if utils.type_string(arg0) == 'HalfTensor':
+            raise NotImplementedError('Cannot call in-place method ' +
+                                      '{} on fp16 Tensors.'.format(fn))
+        else:
+            cast_fn = utils.verbosify(utils.maybe_float, fn, verbose)
+            new_args = utils.casted_args(cast_fn, args, kwargs)
+            return orig_fn(arg0, *new_args, **kwargs)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+# Current RNN approach:
+# - Wrap top-level `RNN` function in thnn backend
+# - Will call into either CudnnRNN or AutogradRNN
+#  - Each of these are factory functions that return a per-iter
+#    `forward` function
+# - We interpose on the factory function to:
+#   1) Interpose on the actual forward function and put in casts
+#   2) Insert an fp16 `flat_weight` if necessary
+def rnn_cast(backend, fn, handle, verbose=False):
+    orig_rnn = utils.get_func(backend, fn)
+    @functools.wraps(orig_rnn)
+    def rnn_wrapper(*args, **kwargs):
+        flat_weight = kwargs.get('flat_weight')
+        if flat_weight is not None:
+            # We replace `flat_weight` with an uninitialized fp16
+            # Tensor. The "actual" weight tensors (provided in `forward`),
+            # will then be set up as ptrs into the buffer and have the
+            # corresponding fp32 values copied in.
+            # We need to call `copy` on the "actual" weights so that the
+            # autograd graph correctly backprops from the wgrads computed
+            # inside cuDNN (on fp16 weights) into the fp32 weights.
+            assert utils.type_string(flat_weight) == 'FloatTensor'
+            if compat.tensor_is_float_tensor() or compat.tensor_is_variable():
+                # Pre-0.4. A little slower, since it zeros out memory.
+                flat_weight_fp16 = flat_weight.new().half().resize_(flat_weight.shape)
+            else:
+                flat_weight_fp16 = torch.empty_like(flat_weight,
+                                                    dtype=torch.float16)
+            kwargs['flat_weight'] = flat_weight_fp16
+        else:
+            flat_weight_fp16 = None
+
+        forward = orig_rnn(*args, **kwargs)
+        @functools.wraps(forward)
+        def fwd_wrapper(*fargs, **fkwargs):
+            assert len(fargs) == 3 or len(fargs) == 4
+            inputs, weights, hiddens = fargs[:3]
+            assert utils.is_fp_tensor(inputs)
+            assert isinstance(weights, list)
+            cast_fn = utils.verbosify(utils.maybe_half,
+                                      fn,
+                                      verbose)
+            new_args = []
+
+            # 0) Inputs
+            new_args.append(cast_fn(inputs))
+
+            # 1) Weights
+            if flat_weight_fp16 is not None:
+                fp16_weights = utils.synthesize_flattened_rnn_weights(
+                    weights, flat_weight_fp16, fn, verbose)
+            else:
+                fp16_weights = [[cast_fn(w) for w in layer]
+                                for layer in weights]
+            new_args.append(fp16_weights)
+
+            # 2) Inputs: either a tuple (for LSTM) or single tensor
+            if isinstance(hiddens, tuple):
+                new_args.append(tuple(cast_fn(x) for x in hiddens))
+            elif utils.is_fp_tensor(hiddens):
+                new_args.append(cast_fn(hiddens))
+            else:
+                # Hiddens can, in principle, be `None` -- pass through
+                new_args.append(hiddens)
+
+            # 3) Batch sizes (0.4 or later only)
+            if len(fargs) == 4:
+                new_args.append(fargs[3])
+
+            return forward(*new_args, **fkwargs)
+        return fwd_wrapper
+    utils.set_func_save(handle, backend, fn, rnn_wrapper)
+
+def new_rnn_cast(fn, handle, verbose=False):
+    # Forward+backward compatibility around https://github.com/pytorch/pytorch/pull/15744
+    # For rnn backend calls that route through _rnn_impls, we must patch the ref
+    # that _rnn_impls stashed.  For rnn backend calls that directly invoke
+    # _VF.<backend>, e.g. _VF.lstm, we can patch onto VariableFunctionsShim,
+    # which in turn has patched the ref named "_VF" in torch.nn.modules.rnn.
+    if utils.has_func(torch.nn.modules.rnn._rnn_impls, fn):
+        mod = torch.nn.modules.rnn._rnn_impls
+    else:
+        mod = torch.nn.modules.rnn._VF
+        assert isinstance(mod, rnn_compat.VariableFunctionsShim)
+        fn = fn.lower()
+    orig_fn = utils.get_func(mod, fn)
+    cast_fn = utils.verbosify(utils.maybe_half, fn, verbose)
+    @functools.wraps(orig_fn)
+    def wrapper(*args, **kwargs):
+        # Exact call signature from modules/rnn.py
+        assert len(args) == 9
+        assert len(kwargs) == 0
+
+        if not _amp_state.handle.is_active():
+            return orig_fn(*args, **kwargs)
+
+        if isinstance(args[6], bool):
+            params_idx = 2 # Not PackedSequence case
+        else:
+            params_idx = 3 # PackedSequence case
+
+        new_args = []
+        for i, arg in enumerate(args):
+            if i == params_idx:
+                num_params = sum([x.numel() for x in arg])
+                fp16_weight_buf = args[0].new_empty((num_params,),
+                                                    dtype=torch.half)
+                casted_weights = utils.new_synthesize_flattened_rnn_weights(
+                    arg, fp16_weight_buf, fn, verbose)
+                new_args.append(casted_weights)
+            elif utils.is_fp_tensor(arg):
+                new_args.append(cast_fn(arg))
+            else:
+                new_args.append(arg)
+
+        return orig_fn(*new_args)
+    utils.set_func_save(handle, mod, fn, wrapper)
+
+def disable_casts(mod, fn, handle):
+    if not utils.has_func(mod, fn):
+        return
+
+    orig_fn = utils.get_func(mod, fn)
+    @functools.wraps(orig_fn)
+    def wrapper(*args, **kwargs):
+        with handle._disable_casts():
+            return orig_fn(*args, **kwargs)
+    utils.set_func_save(handle, mod, fn, wrapper)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/__init__.py
new file mode 100644
index 0000000000..996dbf148f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/__init__.py
@@ -0,0 +1 @@
+from .bottleneck import Bottleneck, SpatialBottleneck
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/bottleneck.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/bottleneck.py
new file mode 100644
index 0000000000..4f18692320
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/bottleneck.py
@@ -0,0 +1,512 @@
+import torch
+import torch.distributed as dist
+from torch import nn
+import fast_bottleneck
+
+def kaiming_uniform_(tensor, a=0, mode='fan_in', nonlinearity='leaky_relu'):
+    weight_tensor_nchw = tensor
+    nn.init.kaiming_uniform_(weight_tensor_nchw, a=a, mode=mode, nonlinearity=nonlinearity)
+
+class FrozenBatchNorm2d(torch.nn.Module):
+    """
+    BatchNorm2d where the batch statistics and the affine parameters are fixed
+    """
+    def __init__(self, n):
+        super(FrozenBatchNorm2d, self).__init__()
+        self.register_buffer("weight", torch.ones(n))
+        self.register_buffer("bias", torch.zeros(n))
+        self.register_buffer("running_mean", torch.zeros(n))
+        self.register_buffer("running_var", torch.ones(n))
+
+    def get_scale_bias(self, nhwc=False):
+        scale = self.weight * self.running_var.rsqrt()
+        bias = self.bias - self.running_mean * scale
+        if nhwc:
+            scale = scale.reshape(1, 1, 1, -1)
+            bias = bias.reshape(1, 1, 1, -1)
+        else:
+            scale = scale.reshape(1, -1, 1, 1)
+            bias = bias.reshape(1, -1, 1, 1)
+        return scale, bias
+
+    def forward(self, x):
+        scale, bias = self.get_scale_bias()
+        return x * scale + bias
+
+
+@torch.jit.script
+def drelu_dscale1(grad_o, output, scale1):
+    relu_mask = (output>0).half()
+    dx_relu = relu_mask * grad_o
+    g1 = dx_relu * scale1
+    return g1, dx_relu
+
+@torch.jit.script
+def drelu_dscale2(grad_o, output, scale1, scale2):
+    relu_mask = (output>0).half()
+    dx_relu = relu_mask * grad_o
+    g1 = dx_relu * scale1
+    g2 = dx_relu * scale2
+    return g1, g2
+
+class BottleneckFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, nhwc, stride_1x1, scale, bias, x, *conv):
+        # TODO: clean up order of tensors
+        args = [x, *conv[0:3], *scale[0:3], *bias[0:3]]
+        ctx.downsample = len(conv) > 3
+        if ctx.downsample:
+            args.append(conv[3])
+            args.append(scale[3])
+            args.append(bias[3])
+
+        # weight buffers are always in nhwc while shape can be nhwc or channels_last
+        # here we pass in flag and let c++ handle it
+        # alternatively, we can put all sizes into a fixed format and pass it in
+        outputs = fast_bottleneck.forward(nhwc, stride_1x1, args)
+        ctx.save_for_backward(*(args+outputs))
+        # save relu outputs for drelu
+        ctx.nhwc = nhwc
+        ctx.stride_1x1 = stride_1x1
+        return outputs[2]
+
+    # backward relu is not exposed, MUL with mask used now
+    # only support dgrad
+    @staticmethod
+    def backward(ctx, grad_o):
+        outputs = ctx.saved_tensors[-3:]
+
+        if ctx.downsample:
+            grad_conv3, grad_conv4 = drelu_dscale2(grad_o, outputs[2], ctx.saved_tensors[6], ctx.saved_tensors[11])
+        else:
+            grad_conv3, grad_conv4 = drelu_dscale1(grad_o, outputs[2], ctx.saved_tensors[6])
+
+        # create input vector for backward
+        t_list = [*ctx.saved_tensors[0:10]]
+        t_list.append(grad_conv3)
+        t_list.append(grad_conv4)
+
+        # outputs used for wgrad and generating drelu mask
+        t_list.append(outputs[0])
+        t_list.append(outputs[1])
+
+        # in case there is downsample
+        if ctx.downsample:
+            t_list.append(ctx.saved_tensors[10])
+
+        grads = fast_bottleneck.backward(ctx.nhwc, ctx.stride_1x1, t_list)
+
+        return (None, None, None, None, *grads)
+
+bottleneck_function = BottleneckFunction.apply
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+    """3x3 convolution with padding"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+                     padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+def conv1x1(in_planes, out_planes, stride=1):
+    """1x1 convolution"""
+    return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+class Bottleneck(torch.nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+    # here we put it at 1x1
+
+    def __init__(self, in_channels, bottleneck_channels, out_channels, stride=1, groups=1,
+                 dilation=1, norm_func=None, use_cudnn=False, explicit_nhwc=False):
+        super(Bottleneck, self).__init__()
+        if groups != 1:
+            raise RuntimeError('Only support groups == 1')
+        if dilation != 1:
+            raise RuntimeError('Only support dilation == 1')
+        if norm_func == None:
+            norm_func = FrozenBatchNorm2d
+        else:
+            raise RuntimeError('Only support frozen BN now.')
+
+        if stride != 1 or in_channels != out_channels:
+            self.downsample = nn.Sequential(
+                conv1x1(in_channels, out_channels, stride),
+                norm_func(out_channels),
+            )
+        else:
+            self.downsample = None
+
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(in_channels, bottleneck_channels, stride)
+        self.conv2 = conv3x3(bottleneck_channels, bottleneck_channels)
+        self.conv3 = conv1x1(bottleneck_channels, out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.stride = stride
+
+        self.bn1 = norm_func(bottleneck_channels)
+        self.bn2 = norm_func(bottleneck_channels)
+        self.bn3 = norm_func(out_channels)
+
+        self.use_cudnn = use_cudnn
+
+        # setup conv weights
+        self.w_conv = [self.conv1.weight, self.conv2.weight, self.conv3.weight]
+        if self.downsample is not None:
+            self.w_conv.append(self.downsample[0].weight)
+
+        # init weight in nchw format before possible transpose
+        for w in self.w_conv:
+            kaiming_uniform_(w, a=1)
+
+        # TODO: prevent unsupported case usage
+        # support cases
+        #                 native      cudnn
+        # normal             yes         no
+        # channel_last       yes        yes
+        # explicit_nhwc       no        yes
+        self.explicit_nhwc = explicit_nhwc
+        if self.explicit_nhwc:
+            for p in self.parameters():
+                with torch.no_grad():
+                    p.data = p.data.permute(0,2,3,1).contiguous()
+        return
+
+    def forward(self, x):
+        if self.use_cudnn:
+            # calculate scale/bias from registered buffers
+            # TODO: make this better
+            s1, b1 = self.bn1.get_scale_bias(self.explicit_nhwc)
+            s2, b2 = self.bn2.get_scale_bias(self.explicit_nhwc)
+            s3, b3 = self.bn3.get_scale_bias(self.explicit_nhwc)
+            w_scale = [s1, s2, s3]
+            w_bias = [b1, b2, b3]
+            if self.downsample is not None:
+                s4, b4 = self.downsample[1].get_scale_bias(self.explicit_nhwc)
+                w_scale.append(s4)
+                w_bias.append(b4)
+
+            out = bottleneck_function(self.explicit_nhwc, self.stride, w_scale, w_bias, x, *self.w_conv)
+            return out
+
+        if self.explicit_nhwc:
+            raise RuntimeError('explicit nhwc with native ops is not supported.')
+
+        # fallback to native ops
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
+
+
+class SpatialBottleneckFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, spatial_group_size, local_rank, comm, stream1, nhwc, stride_1x1, scale, bias, x, *conv):
+        # TODO: clean up order of tensors
+        args = [x, *conv[0:3], *scale[0:3], *bias[0:3]]
+        ctx.downsample = len(conv) > 3
+        if ctx.downsample:
+            args.append(conv[3])
+            args.append(scale[3])
+            args.append(bias[3])
+
+        # weight buffers are always in nhwc while shape can be nhwc or channels_last
+        # here we pass in flag and let c++ handle it
+        # alternatively, we can put all sizes into a fixed format and pass it in
+        outputs = fast_bottleneck.forward_init(nhwc, stride_1x1, args)
+        fast_bottleneck.forward_out1(nhwc, stride_1x1, args, outputs)
+
+        fast_bottleneck.forward_out2(nhwc, stride_1x1, args, outputs)
+
+        # do halo exchange for outputs[0] (out1)
+        # compute halo cells for outputs[1]
+        if spatial_group_size > 1:
+            out1 = outputs[0]
+            N,Hs,W,C = list(out1.shape)
+            stream1.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(stream1):
+                # copy halos to send buffer
+                send_halos = torch.empty((N,2,W,C),dtype=out1.dtype,device=out1.device)
+                send_halos[:,:1,:,:].copy_(out1[:,:1,:,:])
+                send_halos[:,1:,:,:].copy_(out1[:,Hs-1:,:,:])
+                all_halos = torch.empty((N,2*spatial_group_size,W,C),dtype=out1.dtype,device=out1.device)
+                all_halos = [all_halos[:,i*2:(i+1)*2,:,:] for i in range(spatial_group_size)]
+                dist.all_gather(all_halos,send_halos,group=comm)
+                fat_halo = torch.empty((N,3,W,C),dtype=out1.dtype,device=out1.device)
+                top_out1_halo = all_halos[(spatial_group_size+local_rank-1)%spatial_group_size][:,1:,:,:]
+                if local_rank > 0:
+                    fat_halo[:,:1,:,:].copy_(top_out1_halo)
+                    fat_halo[:,1:3,:,:].copy_(out1[:,:2,:,:])
+                    top_out2 = fast_bottleneck.forward_out2_halo(nhwc, fat_halo, args)
+                btm_out1_halo = all_halos[(local_rank+1)%spatial_group_size][:,:1,:,:]
+                if local_rank < spatial_group_size-1:
+                    fat_halo[:,0:2,:,:].copy_(out1[:,Hs-2:,:,:])
+                    fat_halo[:,2:,:,:].copy_(btm_out1_halo)
+                    btm_out2 = fast_bottleneck.forward_out2_halo(nhwc, fat_halo, args)
+            torch.cuda.current_stream().wait_stream(stream1)
+            out2 = outputs[1]
+            if local_rank > 0:
+                out2[:,:1,:,:].copy_(top_out2)
+            if local_rank < spatial_group_size-1:
+                out2[:,Hs-1:,:,:].copy_(btm_out2)
+            
+        fast_bottleneck.forward_rest(nhwc, stride_1x1, args, outputs)
+        # save halos for backward pass
+        if spatial_group_size > 1:
+            ctx.save_for_backward(*(args+outputs+[top_out1_halo,btm_out1_halo]))
+        else:
+            ctx.save_for_backward(*(args+outputs))
+        # save relu outputs for drelu
+        ctx.nhwc = nhwc
+        ctx.stride_1x1 = stride_1x1
+        ctx.spatial_group_size = spatial_group_size
+        ctx.local_rank = local_rank
+        ctx.comm = comm
+        ctx.stream1 = stream1
+        return outputs[2]
+
+    # backward relu is not exposed, MUL with mask used now
+    # only support dgrad
+    @staticmethod
+    def backward(ctx, grad_o):
+        if ctx.spatial_group_size > 1:
+            top_out1_halo = ctx.saved_tensors[-2]
+            btm_out1_halo = ctx.saved_tensors[-1]
+            outputs = ctx.saved_tensors[-5:-2]
+        else:
+            outputs = ctx.saved_tensors[-3:]
+
+        if ctx.downsample:
+            grad_conv3, grad_conv4 = drelu_dscale2(grad_o, outputs[2], ctx.saved_tensors[6], ctx.saved_tensors[11])
+        else:
+            grad_conv3, grad_conv4 = drelu_dscale1(grad_o, outputs[2], ctx.saved_tensors[6])
+
+        # create input vector for backward
+        t_list = [*ctx.saved_tensors[0:10]]
+        t_list.append(grad_conv3)
+        t_list.append(grad_conv4)
+
+        # outputs used for wgrad and generating drelu mask
+        t_list.append(outputs[0])
+        t_list.append(outputs[1])
+
+        # in case there is downsample
+        if ctx.downsample:
+            t_list.append(ctx.saved_tensors[10])
+
+        grads = fast_bottleneck.backward_init(ctx.nhwc, ctx.stride_1x1, t_list)
+        grad_out2 = fast_bottleneck.backward_grad_out2(ctx.nhwc, ctx.stride_1x1, t_list, grads)
+
+        # compute wgrad2 for internal cells
+        wgrad2 = fast_bottleneck.backward_wgrad2(ctx.nhwc, ctx.stride_1x1, t_list, grads, grad_out2)
+
+        # apply wgrad2 halos
+        if ctx.spatial_group_size > 1:
+            if ctx.local_rank > 0:
+                top_grad2_halo = grad_out2[:,:1,:,:]
+                top_wgrad2_halo = fast_bottleneck.backward_wgrad2_halo(ctx.nhwc, ctx.stride_1x1, t_list, grads, top_out1_halo, top_grad2_halo)
+                wgrad2[:,:1,:,:].add_(top_wgrad2_halo)
+            if ctx.local_rank < ctx.spatial_group_size-1:
+                btm_grad2_halo = grad_out2[:,-1:,:,:]
+                btm_wgrad2_halo = fast_bottleneck.backward_wgrad2_halo(ctx.nhwc, ctx.stride_1x1, t_list, grads, btm_out1_halo, btm_grad2_halo)
+                wgrad2[:,-1:,:,:].add_(btm_wgrad2_halo)
+
+        # do halo exchange of grad_out2 here
+        # compute halo cells for grad_out1
+        if ctx.spatial_group_size > 1:
+            N,Hs,W,C = list(grad_out2.shape)
+            ctx.stream1.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(ctx.stream1):
+                # copy halos to send buffer
+                send_halos = torch.empty((N,2,W,C),dtype=grad_out2.dtype,device=grad_out2.device)
+                send_halos[:,:1,:,:].copy_(grad_out2[:,:1,:,:])
+                send_halos[:,1:,:,:].copy_(grad_out2[:,Hs-1:,:,:])
+                all_halos = torch.empty((N,2*ctx.spatial_group_size,W,C),dtype=grad_out2.dtype,device=grad_out2.device)
+                all_halos = [all_halos[:,i*2:(i+1)*2,:,:] for i in range(ctx.spatial_group_size)]
+                dist.all_gather(all_halos,send_halos,group=ctx.comm)
+                relu1 = t_list[12]
+                fat_halo = torch.empty((N,3,W,C),dtype=grad_out2.dtype,device=grad_out2.device)
+                relu_halo = torch.empty((N,3,W,C),dtype=grad_out2.dtype,device=grad_out2.device)
+                if ctx.local_rank > 0:
+                    top_halo = all_halos[ctx.local_rank-1][:,1:,:,:]
+                    fat_halo[:,:1,:,:].copy_(top_halo)
+                    fat_halo[:,1:,:,:].copy_(grad_out2[:,:2,:,:])
+                    relu_halo[:,:1,:,:].zero_()
+                    relu_halo[:,1:,:,:].copy_(relu1[:,:2,:,:])
+                    top_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo(ctx.nhwc, ctx.stride_1x1, t_list, grads, fat_halo, relu_halo)
+                    top_grad_out1_halo = top_grad_out1_halo[:,1:2,:,:]
+                if ctx.local_rank < ctx.spatial_group_size-1:
+                    btm_halo = all_halos[ctx.local_rank+1][:,:1,:,:]
+                    fat_halo[:,:2,:,:].copy_(grad_out2[:,Hs-2:,:,:])
+                    fat_halo[:,2:,:,:].copy_(btm_halo)
+                    relu_halo[:,:2,:,:].copy_(relu1[:,Hs-2:,:,:])
+                    relu_halo[:,2:,:,:].zero_()
+                    btm_grad_out1_halo = fast_bottleneck.backward_grad_out1_halo(ctx.nhwc, ctx.stride_1x1, t_list, grads, fat_halo, relu_halo)
+                    btm_grad_out1_halo = btm_grad_out1_halo[:,1:2,:,:]
+
+        # compute grad_out1 for internal cells
+        grad_out1 = fast_bottleneck.backward_grad_out1(ctx.nhwc, ctx.stride_1x1, t_list, grads, grad_out2)
+
+        # apply halo cells to grad_out1
+        if ctx.spatial_group_size > 1:
+            w = t_list[2]
+            z = t_list[4]
+            relu1 = t_list[12]
+            #print("w.shape = %s, z.shape = %s, relu1.shape = %s" % (str(list(w.shape)), str(list(z.shape)), str(list(relu1.shape))))
+            torch.cuda.current_stream().wait_stream(ctx.stream1)
+            if ctx.local_rank > 0:
+                grad_out1[:,:1,:,:].copy_(top_grad_out1_halo)
+                #print("ctx.local_rank = %d, apply grad_out1 top halo (grad_out1.shape = %s)" % (ctx.local_rank, str(list(grad_out1.shape))))
+            if ctx.local_rank < ctx.spatial_group_size-1:
+                grad_out1[:,Hs-1:,:,:].copy_(btm_grad_out1_halo)
+                #print("ctx.local_rank = %d, apply grad_out1 btm halo (grad_out1.shape = %s)" % (ctx.local_rank, str(list(grad_out1.shape))))
+
+        fast_bottleneck.backward_rest(ctx.nhwc, ctx.stride_1x1, t_list, grads, grad_out2, grad_out1, wgrad2)
+
+        return (None, None, None, None, None, None, None, None, *grads)
+
+spatial_bottleneck_function = SpatialBottleneckFunction.apply
+
+class SpatialBottleneck(torch.nn.Module):
+    # Bottleneck in torchvision places the stride for downsampling at 3x3 convolution(self.conv2)
+    # while original implementation places the stride at the first 1x1 convolution(self.conv1)
+    # according to "Deep residual learning for image recognition"https://arxiv.org/abs/1512.03385.
+    # This variant is also known as ResNet V1.5 and improves accuracy according to
+    # https://ngc.nvidia.com/catalog/model-scripts/nvidia:resnet_50_v1_5_for_pytorch.
+    # here we put it at 1x1
+
+    def __init__(self, in_channels, bottleneck_channels, out_channels, stride=1, groups=1,
+                 dilation=1, norm_func=None, use_cudnn=False, explicit_nhwc=False, 
+                 spatial_group_size=1, communicator=None):
+        super(SpatialBottleneck, self).__init__()
+        if groups != 1:
+            raise RuntimeError('Only support groups == 1')
+        if dilation != 1:
+            raise RuntimeError('Only support dilation == 1')
+        if norm_func == None:
+            norm_func = FrozenBatchNorm2d
+        else:
+            raise RuntimeError('Only support frozen BN now.')
+
+        if stride != 1 or in_channels != out_channels:
+            self.downsample = nn.Sequential(
+                conv1x1(in_channels, out_channels, stride),
+                norm_func(out_channels),
+            )
+        else:
+            self.downsample = None
+
+        # Both self.conv2 and self.downsample layers downsample the input when stride != 1
+        self.conv1 = conv1x1(in_channels, bottleneck_channels, stride)
+        self.conv2 = conv3x3(bottleneck_channels, bottleneck_channels)
+        self.conv3 = conv1x1(bottleneck_channels, out_channels)
+        self.relu = nn.ReLU(inplace=True)
+        self.stride = stride
+
+        self.bn1 = norm_func(bottleneck_channels)
+        self.bn2 = norm_func(bottleneck_channels)
+        self.bn3 = norm_func(out_channels)
+
+        self.use_cudnn = use_cudnn
+
+        # setup conv weights
+        self.w_conv = [self.conv1.weight, self.conv2.weight, self.conv3.weight]
+        if self.downsample is not None:
+            self.w_conv.append(self.downsample[0].weight)
+
+        # init weight in nchw format before possible transpose
+        for w in self.w_conv:
+            kaiming_uniform_(w, a=1)
+
+        # TODO: prevent unsupported case usage
+        # support cases
+        #                 native      cudnn
+        # normal             yes         no
+        # channel_last       yes        yes
+        # explicit_nhwc       no        yes
+        self.explicit_nhwc = explicit_nhwc
+        if self.explicit_nhwc:
+            for p in self.parameters():
+                with torch.no_grad():
+                    p.data = p.data.permute(0,2,3,1).contiguous()
+
+        # spatial communicator
+        self.spatial_group_size = spatial_group_size
+        if spatial_group_size > 1:
+            world_size = dist.get_world_size()
+            num_groups = world_size // spatial_group_size
+            assert(num_groups*spatial_group_size == world_size), "torch.distributed.get_world_size() must be multiple of group_size"
+            rank = dist.get_rank()
+            self.local_rank = rank % spatial_group_size
+            if communicator is None:
+                for group in range(num_groups):
+                    ranks = list(range(group*spatial_group_size,(group+1)*spatial_group_size))
+                    comm = torch.distributed.new_group(ranks=ranks)
+                    if rank in ranks:
+                        self.communicator = comm
+            else:
+                self.communicator = communicator
+            self.stream1 = torch.cuda.Stream()
+            self.spatial_args = self.spatial_group_size, self.local_rank, self.communicator, self.stream1
+        else:
+            self.spatial_args = 1, 0, None, None
+
+        return
+
+    def forward(self, x):
+        if self.use_cudnn:
+            # calculate scale/bias from registered buffers
+            # TODO: make this better
+            s1, b1 = self.bn1.get_scale_bias(self.explicit_nhwc)
+            s2, b2 = self.bn2.get_scale_bias(self.explicit_nhwc)
+            s3, b3 = self.bn3.get_scale_bias(self.explicit_nhwc)
+            w_scale = [s1, s2, s3]
+            w_bias = [b1, b2, b3]
+            if self.downsample is not None:
+                s4, b4 = self.downsample[1].get_scale_bias(self.explicit_nhwc)
+                w_scale.append(s4)
+                w_bias.append(b4)
+
+            out = spatial_bottleneck_function(*self.spatial_args, self.explicit_nhwc, self.stride, w_scale, w_bias, x, *self.w_conv)
+            return out
+
+        if self.explicit_nhwc:
+            raise RuntimeError('explicit nhwc with native ops is not supported.')
+
+        # fallback to native ops
+        identity = x
+
+        out = self.conv1(x)
+        out = self.bn1(out)
+        out = self.relu(out)
+
+        out = self.conv2(out)
+        out = self.bn2(out)
+        out = self.relu(out)
+
+        out = self.conv3(out)
+        out = self.bn3(out)
+
+        if self.downsample is not None:
+            identity = self.downsample(x)
+
+        out += identity
+        out = self.relu(out)
+
+        return out
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/bottleneck_module_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/bottleneck_module_test.py
new file mode 100644
index 0000000000..38aa228426
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/bottleneck_module_test.py
@@ -0,0 +1,272 @@
+import os
+import torch
+from maskrcnn_benchmark.modeling.backbone.resnet import Bottleneck
+from maskrcnn_benchmark.layers.nhwc import nhwc_to_nchw_transform, nchw_to_nhwc_transform
+from maskrcnn_benchmark.layers.nhwc.batch_norm import FrozenBatchNorm2d_NHWC
+from apex.contrib.bottleneck import Bottleneck as FastBottleneck
+from apex.contrib.bottleneck import SpatialBottleneck
+
+
+def single_module_test(ref, rank, world_size, numtype, device, shape, fast, spatial_group_size, in_channels, bottleneck_channels, out_channels, num_groups, stride_in_1x1, stride, dilation, norm_func, nhwc):
+    # inputs + modules
+    with torch.no_grad():
+        input_shape = [1, in_channels] + list(shape)
+        x = torch.randn(input_shape, dtype=numtype, device=device)
+        if nhwc:
+            x = nchw_to_nhwc_transform(x).contiguous()
+        x.requires_grad = True
+        print(x.shape, x.stride())
+
+        #if spatial_group_size > 1:
+        #    fast = False # hack so fast bottleneck can be run against distributed bottleneck
+        #if spatial_group_size == 1:
+        #    fast = False
+
+        if fast:
+            if spatial_group_size == 1:
+                bottleneck = FastBottleneck(
+                    in_channels=in_channels,
+                    bottleneck_channels=bottleneck_channels,
+                    out_channels=out_channels,
+                    stride=stride,
+                    dilation=dilation,
+                    explicit_nhwc=nhwc,
+                    use_cudnn=True)
+            else:
+                bottleneck = SpatialBottleneck(
+                    in_channels=in_channels,
+                    bottleneck_channels=bottleneck_channels,
+                    out_channels=out_channels,
+                    stride=stride,
+                    dilation=dilation,
+                    explicit_nhwc=nhwc,
+                    use_cudnn=True,
+                    spatial_group_size=spatial_group_size)
+        else:
+            bottleneck = Bottleneck(
+                in_channels,
+                bottleneck_channels,
+                out_channels,
+                num_groups,
+                stride_in_1x1,
+                stride,
+                dilation,
+                norm_func,
+                nhwc,
+                spatial_group_size)
+        bottleneck = bottleneck.to(dtype=numtype,device=device)
+        weights = dict(bottleneck.named_parameters())
+
+        if ref is not None:
+            ref_x, _, ref_weights = ref
+            Hs,H = x.shape[1], ref_x.shape[1]
+            assert(Hs*spatial_group_size == H), "Hs not a multiple of H"
+            ref_x = ref_x[:,rank*Hs:(rank+1)*Hs,:,:]
+            x.copy_(ref_x)
+            assert(len(weights) == len(ref_weights)), "Reference weights and weights don't match"
+            for k in weights.keys():
+                weights[k].copy_(ref_weights[k])
+
+    # forward
+    out = bottleneck(x)
+    
+    # gradient output
+    with torch.no_grad():
+        grad_out = torch.randn_like(out)
+        if ref is not None:
+            _, ref_grad_out, _ = ref
+            Hs,H = grad_out.shape[1], ref_grad_out.shape[1]
+            assert(Hs*spatial_group_size == H), "Hs not a multiple of H"
+            ref_grad_out = ref_grad_out[:,rank*Hs:(rank+1)*Hs,:,:]
+            grad_out.copy_(ref_grad_out)
+
+    # backward
+    out.backward(grad_out)
+
+    with torch.no_grad():
+        dgrad = x.grad.detach()
+        
+        wgrad = {}
+        for n,p in bottleneck.named_parameters():
+            wgrad[n] = p.grad.detach()
+
+    if world_size > 1:
+        if spatial_group_size == 1:
+            # broadcast x, grad_out and weights from rank 0
+            with torch.no_grad():
+                torch.distributed.broadcast(x,0)
+                torch.distributed.broadcast(grad_out,0)
+                for k in weights.keys():
+                    torch.distributed.broadcast(weights[k],0)
+        else:
+            # gather dgrad (x.grad), sum wgrad (weights) and out
+            N,Hs,W,C = dgrad.shape
+            H = Hs * spatial_group_size
+            dgrad_gathered = torch.empty((N,H,W,C),dtype=dgrad.dtype,device=dgrad.device)
+            dgrad_tensors = [dgrad_gathered[:,i*Hs:(i+1)*Hs,:,:] for i in range(spatial_group_size)]
+            torch.distributed.all_gather(dgrad_tensors, dgrad)
+            dgrad = dgrad_gathered
+            N,Hs,W,C = list(out.shape)
+            H = Hs * spatial_group_size
+            out_gathered = torch.empty((N,H,W,C),dtype=dgrad.dtype,device=dgrad.device)
+            out_tensors= [out_gathered[:,i*Hs:(i+1)*Hs,:,:] for i in range(spatial_group_size)]
+            torch.distributed.all_gather(out_tensors, out)
+            out = out_gathered
+            for k in wgrad.keys():
+                w = wgrad[k].to(dtype=torch.float64)
+                torch.distributed.all_reduce(w)
+                wgrad[k].copy_(w.to(dtype=wgrad[k].dtype))
+                #torch.distributed.all_reduce(wgrad[k])
+
+    return x, out, grad_out, weights, dgrad, wgrad
+
+
+def module_tests(rank, world_size, numtype, device, fast, spatial_group_sizes, init_args):
+    r = []
+    for ia in init_args:
+        shape = ia[0:4]
+        args = ia[4:]
+        rr = []
+        ref = None
+        for spatial_group_size in spatial_group_sizes:
+            N,H,W,C = shape
+            H = H//spatial_group_size
+            x, out, grad_out, weights, dgrad, wgrad = single_module_test(ref, rank, world_size, numtype, device, [H,W], fast, spatial_group_size, *args)
+            if ref is None:
+                assert(spatial_group_size == 1), "Wrong reference weights"
+                ref = x, grad_out, weights
+            if rank == 0:
+                rr.append( (out, dgrad, wgrad) )
+            if world_size > 1: torch.distributed.barrier()
+        r.append(rr)
+    return r
+
+
+def main():
+    total_num_gpus = int(os.environ["WORLD_SIZE"]) if "WORLD_SIZE" in os.environ else 1
+    distributed = total_num_gpus > 1
+    ngpus = torch.cuda.device_count()
+
+    if distributed:
+        torch.distributed.init_process_group("nccl")
+        rank, world_size = torch.distributed.get_rank(), torch.distributed.get_world_size()
+        is_master = True if rank == 0 else False
+        local_rank = rank % ngpus
+        torch.cuda.set_device(local_rank)
+        spatial_group_size = total_num_gpus
+    else:
+        rank, local_rank, is_master, world_size, spatial_group_size = 0, 0, True, 1, 1
+
+    torch.use_deterministic_algorithms(True)
+    torch.backends.cudnn.benchmark = False
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+    norm_func = FrozenBatchNorm2d_NHWC
+
+    init_args = [
+        (1, 200, 336, 64, 64, 64, 256, 1, True, 1, 1, norm_func, True),
+        (1, 200, 336, 256, 256, 64, 256, 1, True, 1, 1, norm_func, True),
+        (1, 200, 336, 256, 256, 128, 512, 1, True, 2, 1, norm_func, True),
+        (1, 100, 168, 512, 512, 128, 512, 1, True, 1, 1, norm_func, True),
+        (1, 100, 168, 512, 512, 256, 1024, 1, True, 2, 1, norm_func, True),
+        (1, 50, 84, 1024, 1024, 256, 1024, 1, True, 1, 1, norm_func, True),
+        (1, 50, 84, 1024, 1024, 512, 2048, 1, True, 2, 1, norm_func, True),
+        (1, 25, 42, 2048, 2048, 512, 2048, 1, True, 1, 1, norm_func, True),
+        (1, 336, 200, 64, 64, 64, 256, 1, True, 1, 1, norm_func, True),
+        (1, 336, 200, 256, 256, 64, 256, 1, True, 1, 1, norm_func, True),
+        (1, 336, 200, 256, 256, 128, 512, 1, True, 2, 1, norm_func, True),
+        (1, 168, 100, 512, 512, 128, 512, 1, True, 1, 1, norm_func, True),
+        (1, 168, 100, 512, 512, 256, 1024, 1, True, 2, 1, norm_func, True),
+        (1, 84, 50, 1024, 1024, 256, 1024, 1, True, 1, 1, norm_func, True),
+        (1, 84, 50, 1024, 1024, 512, 2048, 1, True, 2, 1, norm_func, True),
+        (1, 42, 25, 2048, 2048, 512, 2048, 1, True, 1, 1, norm_func, True),
+        ]
+    init_args = init_args[0:1]
+
+    # pad H to account for spatial distribution 
+    padded_init_args = []
+    for ia in init_args:
+        N,H,W,C = ia[0:4]
+        m = spatial_group_size * H // (25 if H < W else 42)
+        H = ((H + m - 1) // m) * m
+        args = tuple( [N,H,W,C] + list(ia[4:]) )
+        padded_init_args.append(args)
+    init_args = padded_init_args
+    if rank == 0:
+        for ia in init_args:
+            print(ia)
+
+    spatial_group_sizes = [1]
+    if spatial_group_size > 1:
+        spatial_group_sizes.append(spatial_group_size)
+
+    numtype, device, fast = torch.float16, 'cuda', True
+    r = module_tests(rank, world_size, numtype, device, fast, spatial_group_sizes, init_args)
+    if world_size > 1: torch.distributed.barrier()
+    if rank == 0:
+        for rr in r:
+            print("***")
+            for out, dgrad, wgrad in rr:
+                gr = [("out",out.norm(p=2,dtype=torch.float64).item())]
+                gr = gr + [("dgrad",dgrad.norm(p=2,dtype=torch.float64).item())]
+                gr = gr + [(k+".wgrad",wgrad[k].norm(p=2,dtype=torch.float64).item()) for k in wgrad.keys()]
+                print(gr)
+            if len(rr) == 2:
+                out1, dgrad1, wgrad1 = rr[0]
+                out2, dgrad2, wgrad2 = rr[1]
+
+                rtol = 1e-1
+                out_atol = out1.abs().max().item() * rtol
+                dgrad_atol = dgrad1.abs().max().item() * rtol
+                wgrad_atol = {}
+                for k in wgrad1.keys():
+                    wgrad_atol[k] = wgrad1[k].abs().max().item() * rtol
+
+                gr = [("out",torch.allclose(out1,out2,rtol,out_atol,equal_nan=True))]
+                gr = gr + [("dgrad",torch.allclose(dgrad1,dgrad2,rtol,dgrad_atol,equal_nan=True))]
+                gr = gr + [(k+".wgrad",torch.allclose(wgrad1[k],wgrad2[k],rtol,wgrad_atol[k],equal_nan=True)) for k in wgrad1.keys()]
+                print(gr)
+
+                gr = [("out",(out1-out2).norm(p=2,dtype=torch.float64).item())]
+                gr = gr + [("dgrad",(dgrad1-dgrad2).norm(p=2,dtype=torch.float64).item())]
+                gr = gr + [(k+".wgrad",(wgrad1[k]-wgrad2[k]).norm(p=2,dtype=torch.float64).item()) for k in wgrad1.keys()]
+                print(gr)
+
+                N,H,W,C = out1.shape
+                Hs = H // spatial_group_size
+                Ht = Hs-2
+                print("out1@%d:%d=%s" % (Ht,H,str(out1[0,Ht,:8,:5])))
+                print("out2@%d:%d=%s" % (Ht,H,str(out2[0,Ht,:8,:5])))
+                Ht = Hs-1
+                print("out1@%d:%d=%s" % (Ht,H,str(out1[0,Ht,:8,:5])))
+                print("out2@%d:%d=%s" % (Ht,H,str(out2[0,Ht,:8,:5])))
+                Ht = Hs
+                print("out1@%d:%d=%s" % (Ht,H,str(out1[0,Ht,:8,:5])))
+                print("out2@%d:%d=%s" % (Ht,H,str(out2[0,Ht,:8,:5])))
+                Ht = Hs+1
+                print("out1@%d:%d=%s" % (Ht,H,str(out1[0,Ht,:8,:5])))
+                print("out2@%d:%d=%s" % (Ht,H,str(out2[0,Ht,:8,:5])))
+
+                N,H,W,C = dgrad1.shape
+                Hs = H // spatial_group_size
+                Ht = Hs-2
+                print("dgrad1@%d:%d=%s" % (Ht,H,str(dgrad1[0,Ht,:8,:5])))
+                print("dgrad2@%d:%d=%s" % (Ht,H,str(dgrad2[0,Ht,:8,:5])))
+                Ht = Hs-1
+                print("dgrad1@%d:%d=%s" % (Ht,H,str(dgrad1[0,Ht,:8,:5])))
+                print("dgrad2@%d:%d=%s" % (Ht,H,str(dgrad2[0,Ht,:8,:5])))
+                Ht = Hs
+                print("dgrad1@%d:%d=%s" % (Ht,H,str(dgrad1[0,Ht,:8,:5])))
+                print("dgrad2@%d:%d=%s" % (Ht,H,str(dgrad2[0,Ht,:8,:5])))
+                Ht = Hs+1
+                print("dgrad1@%d:%d=%s" % (Ht,H,str(dgrad1[0,Ht,:8,:5])))
+                print("dgrad2@%d:%d=%s" % (Ht,H,str(dgrad2[0,Ht,:8,:5])))
+
+
+    if world_size > 1: torch.distributed.barrier()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/test.py
new file mode 100644
index 0000000000..2c3c621302
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/bottleneck/test.py
@@ -0,0 +1,71 @@
+import torch
+from bottleneck import Bottleneck
+torch.manual_seed(23337)
+
+# use True to print layerwise sum for all outputs in reference code path
+DEBUG = False#True
+
+for stride, o_channel in [(1,32), (1,128), (2,32)]:
+    print("testing stride ==", stride, ", in_channel == 32 , out_channel ==", o_channel)
+    a_ = torch.randn(17,32,28,28)
+
+    a = a_.cuda().half().to(memory_format=torch.channels_last).requires_grad_()
+    model = Bottleneck(32,8,o_channel,stride=stride).cuda().half().to(memory_format=torch.channels_last)
+
+    # test model
+    b = model(a)
+    b.mean().backward()
+    d_grad = a.grad.float()
+    a.grad = None
+    torch.cuda.synchronize()
+
+    if DEBUG:
+        print("[DEBUG] ref dx :", d_grad.sum().item())
+        # print wgrad. we don't need to reset since later cpp print before accumulation
+        for i, w in enumerate(model.w_conv):
+            print("[DEBUG] ref wgrad{} :".format(i+1), w.grad.sum().item())
+
+    wgrads = []
+    for w in model.w_conv:
+        wgrads.append(w.grad.float())
+
+    model.use_cudnn = True
+    model.zero_grad()
+    c = model(a)
+    c.mean().backward()
+
+    torch.cuda.synchronize()
+    print("comparing native and channels_last:")
+    print("max error fprop:", (b-c).abs().max().item(), "max elem:", b.abs().max().item())
+    print("max error dgrad:", (d_grad-a.grad.float()).abs().max().item(), "max elem:", d_grad.abs().max().item())
+    for i, (w, wgrad) in enumerate(zip(model.w_conv, wgrads)):
+        print("max error wgrad{}:".format(i+1), (wgrad - w.grad.float()).abs().max().item(), "max elem:", wgrad.abs().max().item())
+
+    nhwc_a = a_.permute(0,2,3,1).contiguous().cuda().half().requires_grad_()
+    nhwc_model = Bottleneck(32,8,o_channel,stride=stride,explicit_nhwc=True, use_cudnn=True).cuda().half()
+    for p,q in zip(model.parameters(), nhwc_model.parameters()):
+        # model's storage is already in nhwc, we clone and assign to explicit nhwc model
+        q.data.copy_(p.data.permute(0,2,3,1).contiguous())
+    for p,q in zip(model.buffers(), nhwc_model.buffers()):
+        q.data.copy_(p.data)
+
+    d = nhwc_model(nhwc_a)
+    d.mean().backward()
+    torch.cuda.synchronize()
+
+    # reset reference to cudnn channels_last permute
+    #c_s = c.storage().tolist()
+    #d_s = d.storage().tolist()
+    #print(max([x-y for x,y in zip(c_s,d_s)]))
+    c = c.contiguous(memory_format=torch.contiguous_format).permute(0,2,3,1).contiguous()
+    d_grad = a.grad.float().permute(0,2,3,1).contiguous()
+    wgrads = []
+    for w in model.w_conv:
+        wgrads.append(w.grad.float().permute(0,2,3,1).contiguous())
+
+    torch.cuda.synchronize()
+    print("comparing nhwc and channels_last:")
+    print("max error fprop:", (d-c).abs().max().item(), "max elem:", c.abs().max().item())
+    print("max error dgrad:", (d_grad-nhwc_a.grad.float()).abs().max().item(), "max elem:", d_grad.abs().max().item())
+    for i, (w, wgrad) in enumerate(zip(nhwc_model.w_conv, wgrads)):
+        print("max error wgrad{}:".format(i+1), (wgrad - w.grad.float()).abs().max().item(), "max elem:", wgrad.abs().max().item())
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/bottleneck/bottleneck.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/bottleneck/bottleneck.cpp
new file mode 100644
index 0000000000..4f9475c5cd
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/bottleneck/bottleneck.cpp
@@ -0,0 +1,2486 @@
+#include <ATen/ATen.h>
+#include <ATen/cudnn/Handle.h>  // for getcudnnhandle
+#include <torch/extension.h>
+#include <torch/torch.h>
+#include <vector>
+#include <cudnn_frontend.h>
+
+#include <iostream>
+
+#ifdef DEBUG
+#define DEBUG_MSG(str) do { std::cout << str << std::endl; } while( false )
+#else
+#define DEBUG_MSG(str) do { } while ( false )
+#endif
+
+#ifdef DEBUG_CUDNN
+#define DEBUG_CUDNN_MSG(buf, str) do { buf << str << std::endl; } while( false )
+#else
+#define DEBUG_CUDNN_MSG(buf, str) do { } while ( false )
+#endif
+
+#define checkCudnnErr(...)                                                        \
+    do {                                                                          \
+        int err = checkCudnnError(__VA_ARGS__, #__VA_ARGS__, __FILE__, __LINE__); \
+        if (err) {                                                                \
+            return;                                                    \
+	}                                                                         \
+    } while (0)
+
+
+int checkCudnnError(cudnnStatus_t code, const char* expr, const char* file, int line) {
+    if (code) {
+        printf("CUDNN error at %s:%d, code=%d (%s) in '%s'\n", file, line, (int)code, cudnnGetErrorString(code), expr);
+        return 1;
+    }
+    return 0;
+}
+
+void checkError(cudaError_t code, char const * func, const char *file, const int line, bool abort = true);
+#define checkCUDAError(val) { checkError((val), #val, __FILE__, __LINE__); }    // in-line regular function
+
+void checkError(cudaError_t code, char const * func, const char *file, const int line, bool abort)
+{
+  if (code != cudaSuccess)
+  {
+    const char * errorMessage = cudaGetErrorString(code);
+    fprintf(stderr, "CUDA error returned from \"%s\" at %s:%d, Error code: %d (%s)\n", func, file, line, code, errorMessage);
+    if (abort){
+      cudaDeviceReset();
+      exit(code);
+    }
+  }
+}
+
+void generateStrides(const int64_t* dimA, int64_t* strideA, int nbDims, cudnnTensorFormat_t filterFormat) {
+    // For INT8x4 and INT8x32 we still compute standard strides here to input
+    // into the cuDNN functions. We will manually scale by resizeFactor in the cpu ref.
+    if (filterFormat == CUDNN_TENSOR_NCHW) {
+        strideA[nbDims - 1] = 1;
+        for (int64_t d = nbDims - 2; d >= 0; d--) {
+            strideA[d] = strideA[d + 1] * dimA[d + 1];
+        }
+    } else {
+        // Here we assume that the format is CUDNN_TENSOR_NHWC
+	strideA[1]          = 1;
+        strideA[nbDims - 1] = strideA[1] * dimA[1];
+        for (int64_t d = nbDims - 2; d >= 2; d--) {
+            strideA[d] = strideA[d + 1] * dimA[d + 1];
+        }
+        strideA[0] = strideA[2] * dimA[2];
+    }
+}
+
+
+int getFwdConvDilatedFilterDim(int filterDim, int dilation) {
+    return ((filterDim - 1) * dilation) + 1;
+}
+
+int getFwdConvPaddedImageDim(int tensorDim, int pad) {
+    return tensorDim + (2 * pad);
+}
+
+int getFwdConvOutputDim(
+    int tensorDim,
+    int pad,
+    int filterDim,
+    int stride,
+    int dilation)
+{
+    int p = (getFwdConvPaddedImageDim(tensorDim, pad) - getFwdConvDilatedFilterDim(filterDim, dilation)) / stride + 1;
+    return (p);
+}
+
+enum {
+    X_TENSOR,
+    Y_TENSOR,
+    W_TENSOR,
+    Z_TENSOR,
+    B_TENSOR,
+    AFTERADD_TENSOR,
+    AFTERBIAS_TENSOR,
+    AFTERCONV_TENSOR,
+    OPTIONAL,
+    AFTEROPT_TENSOR,
+};
+
+using common_conv_descriptors =
+    std::tuple<cudnn_frontend::Tensor, cudnn_frontend::Tensor, cudnn_frontend::Tensor, cudnn_frontend::ConvDesc>;
+
+
+common_conv_descriptors
+create_common_descriptors(int64_t* x_dim_padded,
+                          int64_t* padA,
+                          int64_t* convstrideA,
+                          int64_t* dilationA,
+                          int64_t* w_dim_padded,
+                          int64_t* y_dim_padded,
+                          cudnnDataType_t dataType,
+                          cudnnConvolutionMode_t mode) {
+    const int convDim = 2;
+
+    int64_t strideA_padded[4];
+    int64_t outstrideA_padded[4];
+    int64_t filterstrideA_padded[4];
+
+    generateStrides(w_dim_padded, filterstrideA_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(x_dim_padded, strideA_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(y_dim_padded, outstrideA_padded, 4, CUDNN_TENSOR_NHWC);
+
+    return common_conv_descriptors(cudnn_frontend::TensorBuilder()
+                                       .setDim(4, x_dim_padded)
+                                       .setStrides(4, strideA_padded)
+                                       .setId('x')
+                                       .setAlignment(16)
+                                       .setDataType(dataType)
+                                       .build(),
+                                   cudnn_frontend::TensorBuilder()
+                                       .setDim(4, y_dim_padded)
+                                       .setStrides(4, outstrideA_padded)
+                                       .setId('y')
+                                       .setAlignment(16)
+                                       .setDataType(dataType)
+                                       .build(),
+                                   cudnn_frontend::TensorBuilder()
+                                       .setDim(4, w_dim_padded)
+                                       .setStrides(4, filterstrideA_padded)
+                                       .setId('w')
+                                       .setAlignment(16)
+                                       .setDataType(dataType)
+                                       .build(),
+                                   cudnn_frontend::ConvDescBuilder()
+                                       .setDataType(CUDNN_DATA_FLOAT)
+                                       .setMathMode(mode)
+                                       .setNDims(convDim)
+                                       .setStrides(convDim, convstrideA)
+                                       .setPrePadding(convDim, padA)
+                                       .setPostPadding(convDim, padA)
+                                       .setDilation(convDim, dilationA)
+                                       .build());
+}
+
+using common_convbias_descriptors = std::tuple<cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor,
+                                               cudnn_frontend::Tensor>;
+
+common_convbias_descriptors
+create_conv_bias_add_act_descriptors(int64_t* x_dim_padded,
+                                     int64_t* padA,
+                                     int64_t* convstrideA,
+                                     int64_t* dilationA,
+                                     int64_t* w_dim_padded,
+                                     int64_t* y_dim_padded,
+                                     cudnnDataType_t dataType) {
+    const int convDim = 2;
+
+    int64_t b_dim_padded[4];
+    b_dim_padded[0] = 1;
+    b_dim_padded[1] = y_dim_padded[1];
+    b_dim_padded[2] = 1;
+    b_dim_padded[3] = 1;
+
+    int64_t x_stride_padded[4];
+    int64_t y_stride_padded[4];
+    int64_t w_stride_padded[4];
+    int64_t b_stride_padded[4];
+
+    generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);
+
+    return common_convbias_descriptors(cudnn_frontend::TensorBuilder()
+                                           .setDim(4, x_dim_padded)
+                                           .setStrides(4, x_stride_padded)
+                                           .setId('x')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setId('y')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, w_dim_padded)
+                                           .setStrides(4, w_stride_padded)
+                                           .setId('w')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, b_dim_padded)
+                                           .setStrides(4, b_stride_padded)
+                                           .setId('z')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, b_dim_padded)
+                                           .setStrides(4, b_stride_padded)
+                                           .setId('b')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setVirtual()
+                                           .setId('A')  // after add
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setVirtual()
+                                           .setId('B')  // after bias
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setId('C')  // after conv
+                                           .setAlignment(16)
+                                           .setVirtual()
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setId('i')
+                                           .setAlignment(16)
+                                           .setDataType(dataType)
+                                           .build(),
+                                       cudnn_frontend::TensorBuilder()
+                                           .setDim(4, y_dim_padded)
+                                           .setStrides(4, y_stride_padded)
+                                           .setId('D')  // after optional add
+                                           .setAlignment(16)
+                                           .setVirtual()
+                                           .setDataType(dataType)
+                                           .build());
+}
+
+// tensor descriptors used for dgrad
+enum {
+    X_OR_DX_TENSOR,
+    DY_TENSOR,
+    W_OR_DW_TENSOR,
+    SCALE_TENSOR,
+    RELU_TENSOR,
+    AFTER_DCONV_TENSOR,
+    AFTER_DRELU_TENSOR,
+};
+
+using dconv_descriptors = std::tuple<cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor,
+                                     cudnn_frontend::Tensor>;
+
+dconv_descriptors
+create_dconv_descriptors(int64_t* x_dim_padded,
+                         int64_t* padA,
+                         int64_t* convstrideA,
+                         int64_t* dilationA,
+                         int64_t* w_dim_padded,
+                         int64_t* y_dim_padded,
+                         cudnnDataType_t dataType) {
+    const int convDim = 2;
+
+    int64_t b_dim_padded[4];
+    b_dim_padded[0] = 1;
+    b_dim_padded[1] = x_dim_padded[1];
+    b_dim_padded[2] = 1;
+    b_dim_padded[3] = 1;
+
+    int64_t x_stride_padded[4];
+    int64_t y_stride_padded[4];
+    int64_t w_stride_padded[4];
+    int64_t b_stride_padded[4];
+
+    generateStrides(w_dim_padded, w_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(x_dim_padded, x_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(y_dim_padded, y_stride_padded, 4, CUDNN_TENSOR_NHWC);
+    generateStrides(b_dim_padded, b_stride_padded, 4, CUDNN_TENSOR_NHWC);
+
+    return dconv_descriptors(cudnn_frontend::TensorBuilder()
+                             .setDim(4, x_dim_padded)
+                             .setStrides(4, x_stride_padded)
+                             .setId('x')
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, y_dim_padded)
+                             .setStrides(4, y_stride_padded)
+                             .setId('y')
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, w_dim_padded)
+                             .setStrides(4, w_stride_padded)
+                             .setId('w')
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, b_dim_padded)
+                             .setStrides(4, b_stride_padded)
+                             .setId('s')
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, x_dim_padded)
+                             .setStrides(4, x_stride_padded)
+                             .setId('r')
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, x_dim_padded)
+                             .setStrides(4, x_stride_padded)
+                             .setVirtual()
+                             .setId('A')  // after dconv
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build(),
+                             cudnn_frontend::TensorBuilder()
+                             .setDim(4, x_dim_padded)
+                             .setStrides(4, x_stride_padded)
+                             .setVirtual()
+                             .setId('B')  // after drelu
+                             .setAlignment(16)
+                             .setDataType(dataType)
+                             .build());
+}
+
+// create a cache for plan
+std::unordered_map<std::string, cudnn_frontend::ExecutionPlan> plan_cache;
+
+// TODO: better name
+std::string getConvFusionString(int64_t* x_dim_padded,
+                                int64_t* padA,
+                                int64_t* convstrideA,
+                                int64_t* dilationA,
+                                int64_t* w_dim_padded,
+                                cudnnDataType_t dataType,
+                                std::string fusion_string) {
+
+  for(int i=0;i<4;i++) {
+    fusion_string += 'X';
+    fusion_string += std::to_string(x_dim_padded[i]);
+  }
+  for(int i=0;i<4;i++) {
+    fusion_string += 'W';
+    fusion_string += std::to_string(w_dim_padded[i]);
+  }
+  for(int i=0;i<2;i++) {
+    fusion_string += 'P';
+    fusion_string += std::to_string(padA[i]);
+  }
+  for(int i=0;i<2;i++) {
+    fusion_string += 'S';
+    fusion_string += std::to_string(convstrideA[i]);
+  }
+  for(int i=0;i<2;i++) {
+    fusion_string += 'D';
+    fusion_string += std::to_string(dilationA[i]);
+  }
+  fusion_string += 'T';
+  fusion_string += std::to_string(dataType);
+  return fusion_string;
+}
+
+cudnn_frontend::ExecutionPlan& getOrCreatePlan(cudnnHandle_t handle_,
+                                               std::stringstream& log_buf,
+                                               cudnn_frontend::OperationGraph& opGraph,
+                                               std::string cache_string,
+                                               bool use_heuristic = true){
+  auto it = plan_cache.find(cache_string);
+  if (it != plan_cache.end()) {
+    DEBUG_CUDNN_MSG(log_buf, "Found plan in cache");
+    return it->second;
+  } else {
+    if (use_heuristic){
+      // TODO: confirm which mode to use
+      auto heuristics = cudnn_frontend::EngineHeuristicsBuilder()
+        .setOperationGraph(opGraph)
+        .setHeurMode(CUDNN_HEUR_MODE_INSTANT)
+        .build();
+      // try 3 times for now as WAR for no heuristic training
+      int max_tries = 3, count = 0;
+      auto& engine_configs = heuristics.getEngineConfig(max_tries);
+      while(true) {
+        try {
+          plan_cache.emplace(cache_string, std::move(cudnn_frontend::ExecutionPlanBuilder()
+                                                     .setHandle(handle_)
+                                                     .setEngineConfig(engine_configs[count], opGraph.getTag())
+                                                     .build()));
+          break;
+        } catch (cudnn_frontend::cudnnException e) {
+          if (++count == max_tries) throw e;
+        }
+      }
+    }else{
+    DEBUG_CUDNN_MSG(log_buf, "No plan in cache");
+    // How many engines support this operation graph ?
+    auto total_engines = opGraph.getEngineCount();
+    DEBUG_CUDNN_MSG(log_buf, opGraph.describe() << " has " << total_engines << " engines.");
+    // We have to randomly pick one engine from [0, total_engines)
+    // Selecting "0" by default
+    auto engine = cudnn_frontend::EngineBuilder().setGlobalEngineIdx(0).setOperationGraph(opGraph).build();
+    DEBUG_CUDNN_MSG(log_buf, engine.describe());
+    auto& knobs = engine.getSupportedKnobs();
+    for (auto it = std::begin(knobs); it != std::end(knobs); ++it) {
+      DEBUG_CUDNN_MSG(log_buf, it->describe());
+    }
+    if (knobs.begin() != knobs.end()) {
+      DEBUG_CUDNN_MSG(log_buf, "Updated knob choice");
+      knobs.begin()->setChoice(knobs.begin()->getMinValue() + 1);
+      DEBUG_CUDNN_MSG(log_buf, knobs.begin()->describe());
+    }
+
+    // Createmplacee the requisite engine config
+    auto engine_config = cudnn_frontend::EngineConfigBuilder().setEngine(engine).build();
+    DEBUG_CUDNN_MSG(log_buf, engine_config.describe());
+    plan_cache.emplace(cache_string, std::move(cudnn_frontend::ExecutionPlanBuilder().setHandle(handle_).setEngineConfig(engine_config).build()));
+    }
+
+    return plan_cache.find(cache_string)->second;
+  }
+}
+
+void
+run_conv_scale_bias_add_activation(int64_t* x_dim_padded,
+                                   int64_t* pad,
+                                   int64_t* convstride,
+                                   int64_t* dilation,
+                                   int64_t* w_dim_padded,
+                                   int64_t* y_dim_padded,
+                                   cudnnDataType_t dataType,
+                                   at::Half* devPtrX,
+                                   at::Half* devPtrW,
+                                   at::Half* devPtrY,
+                                   at::Half* devPtrZ,
+                                   at::Half* devPtrB,
+                                   at::Half* devPtrI) {
+    cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+    std::stringstream log_buf;
+    try {
+        int convDim = 2;
+
+        // Creates the necessary tensor descriptors
+        common_convbias_descriptors tensors = create_conv_bias_add_act_descriptors(
+            x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+        DEBUG_CUDNN_MSG(log_buf, std::get<X_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<Y_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<W_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<Z_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<B_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERADD_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERBIAS_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERCONV_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<OPTIONAL>(tensors).describe());
+
+        // Define the add operation
+        auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+                           .setMode(CUDNN_POINTWISE_MUL)
+                           .setMathPrecision(CUDNN_DATA_FLOAT)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+        // Define the bias operation
+        auto biasDesc = cudnn_frontend::PointWiseDescBuilder()
+                            .setMode(CUDNN_POINTWISE_ADD)
+                            .setMathPrecision(CUDNN_DATA_FLOAT)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, biasDesc.describe());
+
+        // optional add
+        auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+                            .setMode(CUDNN_POINTWISE_ADD)
+                            .setMathPrecision(CUDNN_DATA_FLOAT)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+        // Define the activation operation
+        auto actDesc = cudnn_frontend::PointWiseDescBuilder()
+                           .setMode(CUDNN_POINTWISE_RELU_FWD)
+                           .setMathPrecision(CUDNN_DATA_FLOAT)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, actDesc.describe());
+
+        // Define the convolution problem
+        auto convDesc = cudnn_frontend::ConvDescBuilder()
+                            .setDataType(CUDNN_DATA_FLOAT)
+                            .setMathMode(CUDNN_CROSS_CORRELATION)
+                            .setNDims(convDim)
+                            .setStrides(convDim, convstride)
+                            .setPrePadding(convDim, pad)
+                            .setPostPadding(convDim, pad)
+                            .setDilation(convDim, dilation)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+        float alpha  = 1.0f;
+        float beta   = 0.0f;
+
+        // Create a convolution Node
+        auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)
+                           .setxDesc(std::get<X_TENSOR>(tensors))
+                           .setwDesc(std::get<W_TENSOR>(tensors))
+                           .setyDesc(std::get<AFTERCONV_TENSOR>(tensors))
+                           .setcDesc(convDesc)
+                           .setAlpha(alpha)
+                           .setBeta(beta)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+        // Create a Add Node with scaling parameters.
+        auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+                           .setxDesc(conv_op.getOutputTensor())
+                           .setbDesc(std::get<Z_TENSOR>(tensors))
+                           .setyDesc(std::get<AFTERADD_TENSOR>(tensors))
+                           .setpwDesc(scaleDesc)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+        // Create a Bias Node.
+        auto bias_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+                           .setxDesc(scale_op.getOutputTensor())
+                           .setbDesc(std::get<B_TENSOR>(tensors))
+                           .setyDesc(std::get<AFTERBIAS_TENSOR>(tensors))
+                           .setpwDesc(biasDesc)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, bias_op.describe());
+
+        // Create a optional add Node.
+        auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+                           .setxDesc(bias_op.getOutputTensor())
+                           .setbDesc(std::get<OPTIONAL>(tensors))
+                           .setyDesc(std::get<AFTEROPT_TENSOR>(tensors))
+                           .setpwDesc(addDesc)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+
+        // Create an Activation Node.
+        auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setxDesc(devPtrI ? add_op.getOutputTensor() : bias_op.getOutputTensor())
+                          .setyDesc(std::get<Y_TENSOR>(tensors))
+                          .setpwDesc(actDesc)
+                          .build();
+        DEBUG_CUDNN_MSG(log_buf, act_op.describe());
+
+        // Create an Operation Graph. In this case it is convolution add bias activation
+        std::array<cudnn_frontend::Operation const*, 5> ops = {&conv_op, &scale_op, &bias_op, devPtrI ? &add_op : &act_op, &act_op};
+
+        auto opGraph = cudnn_frontend::OperationGraphBuilder()
+          .setHandle(handle_)
+          .setOperationGraph(devPtrI ? ops.size() : 4, ops.data())
+          .build();
+
+        // Create string encoding for plan caching
+        auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+        DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+        auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+        DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+        auto workspace_size = plan.getWorkspaceSize();
+        DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+        void* workspace_ptr = nullptr;
+        auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+        if (workspace_size > 0) {
+          workspace_ptr = workspace_tensor.data_ptr<float>();
+        }
+        void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB, devPtrI};
+        int64_t uids[]    = {'x', 'y', 'w', 'z', 'b', 'i'};
+        auto variantPack  = cudnn_frontend::VariantPackBuilder()
+                               .setWorkspacePointer(workspace_ptr)
+          .setDataPointers(devPtrI ? 6 : 5, data_ptrs)
+          .setUids(devPtrI ? 6 : 5, uids)
+                               .build();
+        DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+        cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+        checkCudnnErr(status);
+        cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error");
+    } catch (cudnn_frontend::cudnnException e) {
+      std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+    }
+}
+
+void
+run_conv_scale_bias(int64_t* x_dim_padded,
+                    int64_t* pad,
+                    int64_t* convstride,
+                    int64_t* dilation,
+                    int64_t* w_dim_padded,
+                    int64_t* y_dim_padded,
+                    cudnnDataType_t dataType,
+                    at::Half* devPtrX,
+                    at::Half* devPtrW,
+                    at::Half* devPtrY,
+                    at::Half* devPtrZ,
+                    at::Half* devPtrB) {
+    cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+    std::stringstream log_buf;
+    try {
+        int convDim = 2;
+
+        // Creates the necessary tensor descriptors
+        common_convbias_descriptors tensors = create_conv_bias_add_act_descriptors(
+            x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+        DEBUG_CUDNN_MSG(log_buf, std::get<X_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<Y_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<W_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<Z_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<B_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERADD_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERBIAS_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTERCONV_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<OPTIONAL>(tensors).describe());
+
+        // Define the add operation
+        auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+          .setMode(CUDNN_POINTWISE_MUL)
+          .setMathPrecision(CUDNN_DATA_FLOAT)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+        // Define the bias operation
+        auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+                            .setMode(CUDNN_POINTWISE_ADD)
+                            .setMathPrecision(CUDNN_DATA_FLOAT)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+        // Define the convolution problem
+        auto convDesc = cudnn_frontend::ConvDescBuilder()
+                            .setDataType(CUDNN_DATA_FLOAT)
+                            .setMathMode(CUDNN_CROSS_CORRELATION)
+                            .setNDims(convDim)
+                            .setStrides(convDim, convstride)
+                            .setPrePadding(convDim, pad)
+                            .setPostPadding(convDim, pad)
+                            .setDilation(convDim, dilation)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+        float alpha  = 1.0f;
+        float beta   = 0.0f;
+
+        // Create a convolution Node
+        auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_FORWARD_DESCRIPTOR)
+                           .setxDesc(std::get<X_TENSOR>(tensors))
+                           .setwDesc(std::get<W_TENSOR>(tensors))
+                           .setyDesc(std::get<AFTERCONV_TENSOR>(tensors))
+                           .setcDesc(convDesc)
+                           .setAlpha(alpha)
+                           .setBeta(beta)
+                           .build();
+        DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+        // Create a Add Node with scaling parameters.
+        auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setxDesc(conv_op.getOutputTensor())
+          .setbDesc(std::get<Z_TENSOR>(tensors))
+          .setyDesc(std::get<AFTERADD_TENSOR>(tensors)) // TODO: change enum to aftermul
+          .setpwDesc(scaleDesc)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+        // Create a Bias Node.
+        auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setxDesc(scale_op.getOutputTensor())
+          .setbDesc(std::get<B_TENSOR>(tensors))
+          .setyDesc(std::get<Y_TENSOR>(tensors))
+          .setpwDesc(addDesc)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+        // Create an Operation Graph. In this case it is convolution add bias activation
+        std::array<cudnn_frontend::Operation const*, 3> ops = {&conv_op, &scale_op, &add_op};
+
+        auto opGraph = cudnn_frontend::OperationGraphBuilder()
+          .setHandle(handle_)
+          .setOperationGraph(ops.size(), ops.data())
+          .build();
+
+        // Create string encoding for plan caching
+        auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+        DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+        auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+        DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+        auto workspace_size = plan.getWorkspaceSize();
+        DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+        void* workspace_ptr = nullptr;
+        auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+        if (workspace_size > 0) {
+          workspace_ptr = workspace_tensor.data_ptr<float>();
+        }
+        void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrB};
+        int64_t uids[]    = {'x', 'y', 'w', 'z', 'b'};
+        auto variantPack  = cudnn_frontend::VariantPackBuilder()
+                               .setWorkspacePointer(workspace_ptr)
+          .setDataPointers(5, data_ptrs)
+          .setUids(5, uids)
+                               .build();
+        DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+        cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+        checkCudnnErr(status);
+        cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error");
+    } catch (cudnn_frontend::cudnnException e) {
+      std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+    }
+}
+
+
+void
+run_dconv_drelu_dscale(int64_t* x_dim_padded,
+                       int64_t* pad,
+                       int64_t* convstride,
+                       int64_t* dilation,
+                       int64_t* w_dim_padded,
+                       int64_t* y_dim_padded,
+                       cudnnDataType_t dataType,
+                       at::Half* devPtrX,
+                       at::Half* devPtrW,
+                       at::Half* devPtrY,
+                       at::Half* devPtrZ,
+                       at::Half* devPtrR) {
+    cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+    std::stringstream log_buf;
+    try {
+        int convDim = 2;
+
+        // Creates the necessary tensor descriptors
+        dconv_descriptors tensors = create_dconv_descriptors(
+            x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+        DEBUG_CUDNN_MSG(log_buf, std::get<X_OR_DX_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<DY_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<W_OR_DW_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<SCALE_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<RELU_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DCONV_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DRELU_TENSOR>(tensors).describe());
+
+        // Define the convolution problem
+        auto convDesc = cudnn_frontend::ConvDescBuilder()
+                            .setDataType(CUDNN_DATA_FLOAT)
+                            .setMathMode(CUDNN_CROSS_CORRELATION)
+                            .setNDims(convDim)
+                            .setStrides(convDim, convstride)
+                            .setPrePadding(convDim, pad)
+                            .setPostPadding(convDim, pad)
+                            .setDilation(convDim, dilation)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+        // Define the activation backward operation
+        auto actDesc = cudnn_frontend::PointWiseDescBuilder()
+          .setMode(CUDNN_POINTWISE_RELU_BWD)
+          .setMathPrecision(CUDNN_DATA_FLOAT)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, actDesc.describe());
+
+        // Define the scale backward operation
+        auto scaleDesc = cudnn_frontend::PointWiseDescBuilder()
+          .setMode(CUDNN_POINTWISE_MUL)
+          .setMathPrecision(CUDNN_DATA_FLOAT)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, scaleDesc.describe());
+
+        float alpha  = 1.0f;
+        float beta   = 0.0f;
+
+        // Create a convolution Node
+        auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)
+          .setdxDesc(std::get<AFTER_DCONV_TENSOR>(tensors))
+          .setwDesc(std::get<W_OR_DW_TENSOR>(tensors))
+          .setdyDesc(std::get<DY_TENSOR>(tensors))
+          .setcDesc(convDesc)
+          .setAlpha(alpha)
+          .setBeta(beta)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+        // TODO: do we need getOutputTensor(), and what it returns in backward case?
+        // Create an relu backward Node.
+        auto act_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setdyDesc(std::get<AFTER_DCONV_TENSOR>(tensors))
+          .setxDesc(std::get<RELU_TENSOR>(tensors))
+          .setdxDesc(std::get<AFTER_DRELU_TENSOR>(tensors))
+          .setpwDesc(actDesc)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, act_op.describe());
+
+        // Create a Scale Node.
+        auto scale_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setxDesc(std::get<AFTER_DRELU_TENSOR>(tensors))
+          .setbDesc(std::get<SCALE_TENSOR>(tensors))
+          .setyDesc(std::get<X_OR_DX_TENSOR>(tensors))
+          .setpwDesc(scaleDesc)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, scale_op.describe());
+
+        // Create an Operation Graph. In this case it is convolution add bias activation
+        std::array<cudnn_frontend::Operation const*, 3> ops = {&conv_op, &act_op, &scale_op};
+
+        auto opGraph = cudnn_frontend::OperationGraphBuilder()
+          .setHandle(handle_)
+          .setOperationGraph(ops.size(), ops.data())
+          .build();
+
+        // Create string encoding for plan caching
+        auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+        DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+        auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+        DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+        auto workspace_size = plan.getWorkspaceSize();
+        DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+        void* workspace_ptr = nullptr;
+        auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+        if (workspace_size > 0) {
+          workspace_ptr = workspace_tensor.data_ptr<float>();
+        }
+        void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrZ, devPtrR};
+        int64_t uids[]    = {'x', 'y', 'w', 's', 'r'};
+        auto variantPack  = cudnn_frontend::VariantPackBuilder()
+          .setWorkspacePointer(workspace_ptr)
+          .setDataPointers(5, data_ptrs)
+          .setUids(5, uids)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+        cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+        checkCudnnErr(status);
+        cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error");
+    } catch (cudnn_frontend::cudnnException e) {
+      std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+    }
+}
+
+void
+run_dconv(int64_t* x_dim_padded,
+          int64_t* pad,
+          int64_t* convstride,
+          int64_t* dilation,
+          int64_t* w_dim_padded,
+          int64_t* y_dim_padded,
+          cudnnDataType_t dataType,
+          at::Half* devPtrX,
+          at::Half* devPtrW,
+          at::Half* devPtrY,
+          cudnnBackendDescriptorType_t mode) {
+    cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+    std::stringstream log_buf;
+    try {
+        int convDim = 2;
+
+        // Creates the necessary tensor descriptors
+        dconv_descriptors tensors = create_dconv_descriptors(
+            x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+        DEBUG_CUDNN_MSG(log_buf, std::get<X_OR_DX_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<DY_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<W_OR_DW_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<SCALE_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<RELU_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DCONV_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DRELU_TENSOR>(tensors).describe());
+
+        // Define the convolution problem
+        auto convDesc = cudnn_frontend::ConvDescBuilder()
+                            .setDataType(CUDNN_DATA_FLOAT)
+                            .setMathMode(CUDNN_CROSS_CORRELATION)
+                            .setNDims(convDim)
+                            .setStrides(convDim, convstride)
+                            .setPrePadding(convDim, pad)
+                            .setPostPadding(convDim, pad)
+                            .setDilation(convDim, dilation)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+        float alpha  = 1.0f;
+        float beta   = 0.0f;
+
+        // Create a convolution Node
+        // mode should be one of following
+        // CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR
+        // CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR
+        auto conv_op_builder = cudnn_frontend::OperationBuilder(mode);
+        if (mode == CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR) {
+          conv_op_builder.setdxDesc(std::get<X_OR_DX_TENSOR>(tensors))
+            .setwDesc(std::get<W_OR_DW_TENSOR>(tensors))
+            .setdyDesc(std::get<DY_TENSOR>(tensors))
+            .setcDesc(convDesc)
+            .setAlpha(alpha)
+            .setBeta(beta);
+        }
+        else {
+          conv_op_builder.setxDesc(std::get<X_OR_DX_TENSOR>(tensors))
+            .setdwDesc(std::get<W_OR_DW_TENSOR>(tensors))
+            .setdyDesc(std::get<DY_TENSOR>(tensors))
+            .setcDesc(convDesc)
+            .setAlpha(alpha)
+            .setBeta(beta);
+        }
+        auto conv_op = conv_op_builder.build();
+        DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+        // Create an Operation Graph. In this case it is convolution add bias activation
+        std::array<cudnn_frontend::Operation const*, 1> ops = {&conv_op};
+
+        auto opGraph = cudnn_frontend::OperationGraphBuilder()
+          .setHandle(handle_)
+          .setOperationGraph(ops.size(), ops.data())
+          .build();
+
+        // Create string encoding for plan caching
+        auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+        DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+        auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+        DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+        auto workspace_size = plan.getWorkspaceSize();
+        DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+        void* workspace_ptr = nullptr;
+        auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+        if (workspace_size > 0) {
+          workspace_ptr = workspace_tensor.data_ptr<float>();
+        }
+        void* data_ptrs[] = {devPtrX, devPtrY, devPtrW};
+        int64_t uids[]    = {'x', 'y', 'w'};
+        auto variantPack  = cudnn_frontend::VariantPackBuilder()
+          .setWorkspacePointer(workspace_ptr)
+          .setDataPointers(3, data_ptrs)
+          .setUids(3, uids)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+        cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+        checkCudnnErr(status);
+        cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error");
+    } catch (cudnn_frontend::cudnnException e) {
+      std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+    }
+}
+
+void
+run_dconv_add(int64_t* x_dim_padded,
+              int64_t* pad,
+              int64_t* convstride,
+              int64_t* dilation,
+              int64_t* w_dim_padded,
+              int64_t* y_dim_padded,
+              cudnnDataType_t dataType,
+              at::Half* devPtrX,
+              at::Half* devPtrW,
+              at::Half* devPtrY,
+              at::Half* devPtrR) {
+    cudnnHandle_t handle_ = torch::native::getCudnnHandle();
+    std::stringstream log_buf;
+    try {
+        int convDim = 2;
+
+        // Creates the necessary tensor descriptors
+        dconv_descriptors tensors = create_dconv_descriptors(
+            x_dim_padded, pad, convstride, dilation, w_dim_padded, y_dim_padded, dataType);
+        DEBUG_CUDNN_MSG(log_buf, std::get<X_OR_DX_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<DY_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<W_OR_DW_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<SCALE_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<RELU_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DCONV_TENSOR>(tensors).describe());
+        DEBUG_CUDNN_MSG(log_buf, std::get<AFTER_DRELU_TENSOR>(tensors).describe());
+
+        // Define the convolution problem
+        auto convDesc = cudnn_frontend::ConvDescBuilder()
+                            .setDataType(CUDNN_DATA_FLOAT)
+                            .setMathMode(CUDNN_CROSS_CORRELATION)
+                            .setNDims(convDim)
+                            .setStrides(convDim, convstride)
+                            .setPrePadding(convDim, pad)
+                            .setPostPadding(convDim, pad)
+                            .setDilation(convDim, dilation)
+                            .build();
+        DEBUG_CUDNN_MSG(log_buf, convDesc.describe());
+
+        // Define the add backward operation
+        auto addDesc = cudnn_frontend::PointWiseDescBuilder()
+          .setMode(CUDNN_POINTWISE_ADD)
+          .setMathPrecision(CUDNN_DATA_FLOAT)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, addDesc.describe());
+
+        float alpha  = 1.0f;
+        float beta   = 0.0f;
+
+        // Create a convolution Node
+        auto conv_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR)
+          .setdxDesc(std::get<AFTER_DCONV_TENSOR>(tensors))
+          .setwDesc(std::get<W_OR_DW_TENSOR>(tensors))
+          .setdyDesc(std::get<DY_TENSOR>(tensors))
+          .setcDesc(convDesc)
+          .setAlpha(alpha)
+          .setBeta(beta)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, conv_op.describe());
+
+        // TODO: do we need getOutputTensor(), and what it returns in backward case?
+        // Create add Node.
+        auto add_op = cudnn_frontend::OperationBuilder(CUDNN_BACKEND_OPERATION_POINTWISE_DESCRIPTOR)
+          .setxDesc(std::get<AFTER_DCONV_TENSOR>(tensors))
+          .setbDesc(std::get<RELU_TENSOR>(tensors))
+          .setyDesc(std::get<X_OR_DX_TENSOR>(tensors))
+          .setpwDesc(addDesc)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, add_op.describe());
+
+        // Create an Operation Graph. In this case it is convolution add bias activation
+        std::array<cudnn_frontend::Operation const*, 2> ops = {&conv_op, &add_op};
+
+        auto opGraph = cudnn_frontend::OperationGraphBuilder()
+          .setHandle(handle_)
+          .setOperationGraph(ops.size(), ops.data())
+          .build();
+
+        // Create string encoding for plan caching
+        auto cache_string = getConvFusionString(x_dim_padded, pad, convstride, dilation, w_dim_padded, dataType, opGraph.getTag());
+        DEBUG_CUDNN_MSG(log_buf, "[convstring] " << cache_string);
+
+        auto& plan = getOrCreatePlan(handle_, log_buf, opGraph, cache_string);
+        DEBUG_CUDNN_MSG(log_buf, "Plan tag: " << plan.getTag());
+
+        auto workspace_size = plan.getWorkspaceSize();
+        DEBUG_CUDNN_MSG(log_buf, plan.describe() << " requires workspace " << workspace_size);
+
+        void* workspace_ptr = nullptr;
+        auto workspace_tensor = at::empty({(workspace_size+3)/4}, at::TensorOptions(at::kCUDA).dtype(at::kFloat));
+        if (workspace_size > 0) {
+          workspace_ptr = workspace_tensor.data_ptr<float>();
+        }
+        void* data_ptrs[] = {devPtrX, devPtrY, devPtrW, devPtrR};
+        int64_t uids[]    = {'x', 'y', 'w', 'r'};
+        auto variantPack  = cudnn_frontend::VariantPackBuilder()
+          .setWorkspacePointer(workspace_ptr)
+          .setDataPointers(4, data_ptrs)
+          .setUids(4, uids)
+          .build();
+        DEBUG_CUDNN_MSG(log_buf, "variantPack " << variantPack.describe());
+        cudnnStatus_t status = cudnnBackendExecute(handle_, plan.get_raw_desc(), variantPack.get_raw_desc());
+        checkCudnnErr(status);
+        cudnn_frontend::throw_if([status]() { return (status != CUDNN_STATUS_SUCCESS); }, "Plan execute error");
+    } catch (cudnn_frontend::cudnnException e) {
+      std::cout << log_buf.str() << "[ERROR] Exception " << e.what() << std::endl;
+    }
+}
+
+
+// inputs contains x,w,z,b,(i)
+std::vector<at::Tensor> bottleneck_forward(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs) {
+
+  std::cout << std::fixed;
+  // create output vector
+  std::vector<at::Tensor> outputs;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // setup dimensions
+  int64_t dimA[]         = {0, 0, 0, 0};
+  int64_t filterdimA1[]  = {0, 0, 0, 0};
+  int64_t filterdimA2[]  = {0, 0, 0, 0};
+  int64_t filterdimA3[]  = {0, 0, 0, 0};
+  int64_t filterdimA4[]  = {0, 0, 0, 0};
+
+  // All dim calculation after this order of n,c,h,w
+  int axis[] {0,1,2,3};
+  if (explicit_nhwc) {
+    axis[0] = 0;
+    axis[1] = 3;
+    axis[2] = 1;
+    axis[3] = 2;
+  }
+  for (int dim=0;dim<4;dim++) {
+    dimA[dim] = inputs[0].size(axis[dim]);
+    filterdimA1[dim] = inputs[1].size(axis[dim]);
+    filterdimA2[dim] = inputs[2].size(axis[dim]);
+    filterdimA3[dim] = inputs[3].size(axis[dim]);
+  }
+  if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+    for (int dim=0;dim<4;dim++) {
+      filterdimA4[dim] = inputs[10].size(axis[dim]);
+    }
+  }
+
+  // output dim in n,c,h,w used by backend
+  int64_t outdimA1[]     = {0, 0, 0, 0}; // Computed Below
+  int64_t outdimA2[]     = {0, 0, 0, 0}; // Computed Below
+  int64_t outdimA3[]     = {0, 0, 0, 0}; // Computed Below
+
+  // use these fixed value for test run
+  int64_t padA[]        = {0, 0};
+  int64_t padA1[]        = {1, 1};
+  int64_t dilationA[] = {1, 1};
+  int64_t convstrideA[] = {1, 1};
+  int64_t convstride1X1[] = {stride_1X1, stride_1X1};
+
+  // compute output from pad/stride/dilation
+  outdimA1[0] = dimA[0];
+  outdimA1[1] = filterdimA1[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA1[dim + 2] = getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);
+  }
+
+  outdimA2[0] = outdimA1[0];
+  outdimA2[1] = filterdimA2[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA2[dim + 2] = getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);
+  }
+
+  outdimA3[0] = outdimA2[0];
+  outdimA3[1] = filterdimA3[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA3[dim + 2] = getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);
+  }
+
+  // Create output tensor in the correct shape in pytorch's view
+  int64_t outdim1[]     = {0, 0, 0, 0};
+  int64_t outdim2[]     = {0, 0, 0, 0};
+  int64_t outdim3[]     = {0, 0, 0, 0};
+  if (explicit_nhwc) {
+    axis[0] = 0;
+    axis[1] = 2;
+    axis[2] = 3;
+    axis[3] = 1;
+  }
+  for (int dim=0;dim<4;dim++) {
+    outdim1[dim] = outdimA1[axis[dim]];
+    outdim2[dim] = outdimA2[axis[dim]];
+    outdim3[dim] = outdimA3[axis[dim]];
+  }
+
+  // run
+  at::Half* x = inputs[0].data_ptr<at::Half>();
+  at::Half* w = inputs[1].data_ptr<at::Half>();
+  at::Half* z = inputs[4].data_ptr<at::Half>();
+  at::Half* b = inputs[7].data_ptr<at::Half>();
+  auto out1 = at::empty(outdim1, inputs[0].type(), output_format);
+  at::Half* y1 = out1.data_ptr<at::Half>();
+
+  run_conv_scale_bias_add_activation(dimA,
+                                     padA,
+                                     convstride1X1,
+                                     dilationA,
+                                     filterdimA1,
+                                     outdimA1,
+                                     CUDNN_DATA_HALF,
+                                     x,
+                                     w,
+                                     y1,
+                                     z,
+                                     b,
+                                     nullptr);
+
+  DEBUG_MSG("[DEBUG] new relu1 : " << out1.to(at::kFloat).sum().item<float>());
+
+  w = inputs[2].data_ptr<at::Half>();
+  z = inputs[5].data_ptr<at::Half>();
+  b = inputs[8].data_ptr<at::Half>();
+  auto out2 = at::empty(outdim2, inputs[0].type(), output_format);
+  at::Half* y2 = out2.data_ptr<at::Half>();
+
+  run_conv_scale_bias_add_activation(outdimA1,
+                                     padA1,
+                                     convstrideA,
+                                     dilationA,
+                                     filterdimA2,
+                                     outdimA2,
+                                     CUDNN_DATA_HALF,
+                                     y1,
+                                     w,
+                                     y2,
+                                     z,
+                                     b,
+                                     nullptr);
+  DEBUG_MSG("[DEBUG] new relu2 : " << out2.to(at::kFloat).sum().item<float>());
+
+  // create output of conv3
+  auto out3 = at::empty(outdim3, inputs[0].type(), output_format);
+  at::Half* y3 = out3.data_ptr<at::Half>();
+
+  // create output of conv4 that may exist
+  auto identity = at::empty_like(out3);
+  at::Half* yi = identity.data_ptr<at::Half>();
+
+  if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]){
+
+    w = inputs[10].data_ptr<at::Half>();
+    z = inputs[11].data_ptr<at::Half>();
+    b = inputs[12].data_ptr<at::Half>();
+    run_conv_scale_bias(dimA,
+                        padA,
+                        convstride1X1,
+                        dilationA,
+                        filterdimA4,
+                        outdimA3,
+                        CUDNN_DATA_HALF,
+                        x,
+                        w,
+                        yi,
+                        z,
+                        b);
+    DEBUG_MSG("[DEBUG] new downsample : " << identity.to(at::kFloat).sum().item<float>());
+  }
+  else {
+    yi = x;
+  }
+
+  w = inputs[3].data_ptr<at::Half>();
+  z = inputs[6].data_ptr<at::Half>();
+  b = inputs[9].data_ptr<at::Half>();
+
+  run_conv_scale_bias_add_activation(outdimA2,
+                                     padA,
+                                     convstrideA,
+                                     dilationA,
+                                     filterdimA3,
+                                     outdimA3,
+                                     CUDNN_DATA_HALF,
+                                     y2,
+                                     w,
+                                     y3,
+                                     z,
+                                     b,
+                                     yi);
+  DEBUG_MSG("[DEBUG] new relu3 : " << out3.to(at::kFloat).sum().item<float>());
+
+  outputs.push_back(out1);
+  outputs.push_back(out2);
+  outputs.push_back(out3);
+
+  return outputs;
+}
+
+std::vector<at::Tensor> bottleneck_backward(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs) {
+
+  bool requires_grad = inputs[0].requires_grad();
+
+  std::cout << std::fixed;
+  // create output vector
+  std::vector<at::Tensor> outputs;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // setup dimensions
+  int64_t dimA[]         = {0, 0, 0, 0};
+  int64_t filterdimA1[]  = {0, 0, 0, 0};
+  int64_t filterdimA2[]  = {0, 0, 0, 0};
+  int64_t filterdimA3[]  = {0, 0, 0, 0};
+  int64_t filterdimA4[]  = {0, 0, 0, 0};
+
+  // All dim calculation after this order of n,c,h,w
+  int axis[] {0,1,2,3};
+  if (explicit_nhwc) {
+    axis[0] = 0;
+    axis[1] = 3;
+    axis[2] = 1;
+    axis[3] = 2;
+  }
+  for (int dim=0;dim<4;dim++) {
+    dimA[dim] = inputs[0].size(axis[dim]);
+    filterdimA1[dim] = inputs[1].size(axis[dim]);
+    filterdimA2[dim] = inputs[2].size(axis[dim]);
+    filterdimA3[dim] = inputs[3].size(axis[dim]);
+  }
+  if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+    for (int dim=0;dim<4;dim++) {
+      filterdimA4[dim] = inputs[14].size(axis[dim]);
+    }
+  }
+
+  // output dim in n,c,h,w used by backend
+  int64_t outdimA1[]     = {0, 0, 0, 0}; // Computed Below
+  int64_t outdimA2[]     = {0, 0, 0, 0}; // Computed Below
+  int64_t outdimA3[]     = {0, 0, 0, 0}; // Computed Below
+
+  // use these fixed value for test run
+  int64_t padA[]        = {0, 0};
+  int64_t padA1[]        = {1, 1};
+  int64_t dilationA[] = {1, 1};
+  int64_t convstrideA[] = {1, 1};
+  int64_t convstride1X1[] = {stride_1X1, stride_1X1};
+
+  // compute output from pad/stride/dilation
+  outdimA1[0] = dimA[0];
+  outdimA1[1] = filterdimA1[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA1[dim + 2] = getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);
+  }
+
+  outdimA2[0] = outdimA1[0];
+  outdimA2[1] = filterdimA2[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA2[dim + 2] = getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);
+  }
+
+  outdimA3[0] = outdimA2[0];
+  outdimA3[1] = filterdimA3[0];
+  for (int dim = 0; dim < 2; dim++) {
+    outdimA3[dim + 2] = getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);
+  }
+
+  // Create output tensor in the correct shape in pytorch's view
+  int64_t outdim1[]     = {0, 0, 0, 0};
+  int64_t outdim2[]     = {0, 0, 0, 0};
+  int64_t outdim3[]     = {0, 0, 0, 0};
+  if (explicit_nhwc) {
+    axis[0] = 0;
+    axis[1] = 2;
+    axis[2] = 3;
+    axis[3] = 1;
+  }
+  for (int dim=0;dim<4;dim++) {
+    outdim1[dim] = outdimA1[axis[dim]];
+    outdim2[dim] = outdimA2[axis[dim]];
+    outdim3[dim] = outdimA3[axis[dim]];
+  }
+
+  // dconv3+drelu2+dscale2
+  at::Half* conv_in = inputs[13].data_ptr<at::Half>();
+  at::Half* dy3 = inputs[10].data_ptr<at::Half>();
+
+  DEBUG_MSG("[DEBUG] new dconv3 : " << inputs[10].to(at::kFloat).sum().item<float>());
+
+  // wgrad
+  auto wgrad3 = at::empty_like(inputs[3]);
+  at::Half* dw3 = wgrad3.data_ptr<at::Half>();
+  run_dconv(outdimA2,
+            padA,
+            convstrideA,
+            dilationA,
+            filterdimA3,
+            outdimA3,
+            CUDNN_DATA_HALF,
+            conv_in,
+            dw3,
+            dy3,
+            CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+  // dgrad
+  auto grad_out2 = at::empty(outdim2, inputs[0].type(), output_format);
+  at::Half* dy2 = grad_out2.data_ptr<at::Half>();
+  at::Half* w = inputs[3].data_ptr<at::Half>();
+  at::Half* z = inputs[5].data_ptr<at::Half>();
+
+  at::Half* relu2 = inputs[13].data_ptr<at::Half>();
+
+  run_dconv_drelu_dscale(outdimA2,
+                         padA,
+                         convstrideA,
+                         dilationA,
+                         filterdimA3,
+                         outdimA3,
+                         CUDNN_DATA_HALF,
+                         dy2,
+                         w,
+                         dy3,
+                         z,
+                         relu2);
+
+  DEBUG_MSG("[DEBUG] new dconv2 : " << grad_out2.to(at::kFloat).sum().item<float>());
+
+  // dconv2+drelu1+dscale1
+  conv_in = inputs[12].data_ptr<at::Half>();
+
+  // wgrad
+  auto wgrad2 = at::empty_like(inputs[2]);
+  at::Half* dw2 = wgrad2.data_ptr<at::Half>();
+  run_dconv(outdimA1,
+            padA1,
+            convstrideA,
+            dilationA,
+            filterdimA2,
+            outdimA2,
+            CUDNN_DATA_HALF,
+            conv_in,
+            dw2,
+            dy2,
+            CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+  // dgrad
+  auto grad_out1 = at::empty(outdim1, inputs[0].type(), output_format);
+  at::Half* dy1 = grad_out1.data_ptr<at::Half>();
+  w = inputs[2].data_ptr<at::Half>();
+  z = inputs[4].data_ptr<at::Half>();
+
+  at::Half* relu1 = inputs[12].data_ptr<at::Half>();
+  // fused dgrad
+  run_dconv_drelu_dscale(outdimA1,
+                         padA1,
+                         convstrideA,
+                         dilationA,
+                         filterdimA2,
+                         outdimA2,
+                         CUDNN_DATA_HALF,
+                         dy1,
+                         w,
+                         dy2,
+                         z,
+                         relu1);
+
+/*
+  // backward strided conv cannot be fused
+  // if stride == 1 but channel changes, we can fuse here
+  if (stride_1X1 != 1){
+    // dgrad
+    run_dconv(outdimA1,
+              padA1,
+              convstride1X1,
+              dilationA,
+              filterdimA2,
+              outdimA2,
+              CUDNN_DATA_HALF,
+              dy1,
+              w,
+              dy2,
+              CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+
+    // mul fused mask
+    grad_out1.mul_(inputs[15]);
+  }
+  else {
+    at::Half* relu1 = inputs[12].data_ptr<at::Half>();
+    // fused dgrad
+    run_dconv_drelu_dscale(outdimA1,
+                           padA1,
+                           convstride1X1,
+                           dilationA,
+                           filterdimA2,
+                           outdimA2,
+                           CUDNN_DATA_HALF,
+                           dy1,
+                           w,
+                           dy2,
+                           z,
+                           relu1);
+  }
+*/
+  DEBUG_MSG("[DEBUG] new dconv1 : " << grad_out1.to(at::kFloat).sum().item<float>());
+
+  // create grads of conv4 that may exist
+  auto grad_x_conv4 = at::empty_like(inputs[0]);
+  at::Half* dx_conv4 = grad_x_conv4.data_ptr<at::Half>();
+  at::Tensor wgrad4;
+
+  // x used for dconv1 and dconv4 wgrad
+  at::Half* x = inputs[0].data_ptr<at::Half>();
+
+  if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]){
+    w = inputs[14].data_ptr<at::Half>();
+    at::Half* dy_conv4 = inputs[11].data_ptr<at::Half>();
+    if (requires_grad) {
+      run_dconv(dimA,
+                padA,
+                convstride1X1,
+                dilationA,
+                filterdimA4,
+                outdimA3,
+                CUDNN_DATA_HALF,
+                dx_conv4,
+                w,
+                dy_conv4,
+                CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+      // we don't print here since we can't hook out this grad in pytorch alone to compare, due to addition with dx
+      // DEBUG_MSG("[DEBUG] new dx_identity : " << grad_x_conv4.to(at::kFloat).sum().item<float>());
+    }
+    // wgrad
+    wgrad4 = at::empty_like(inputs[14]);
+    at::Half* dw4 = wgrad4.data_ptr<at::Half>();
+    run_dconv(dimA,
+              padA,
+              convstride1X1,
+              dilationA,
+              filterdimA4,
+              outdimA3,
+              CUDNN_DATA_HALF,
+              x,
+              dw4,
+              dy_conv4,
+              CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+  }
+  else {
+    // if there is no downsample, dx_conv4 is fork of drelu3
+    dx_conv4 = inputs[11].data_ptr<at::Half>();
+  }
+
+  // dconv1+add
+  // wgrad
+  auto wgrad1 = at::empty_like(inputs[1]);
+  at::Half* dw1 = wgrad1.data_ptr<at::Half>();
+  run_dconv(dimA,
+            padA,
+            convstride1X1,
+            dilationA,
+            filterdimA1,
+            outdimA1,
+            CUDNN_DATA_HALF,
+            x,
+            dw1,
+            dy1,
+            CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+  // dgrad
+  w = inputs[1].data_ptr<at::Half>();
+  auto grad_x = at::empty_like(inputs[0]);
+  at::Half* dx = grad_x.data_ptr<at::Half>();
+
+  // backward strided conv cannot be fused
+  // if stride == 1 but channel changes, we can fuse here
+  if (requires_grad){
+    if (stride_1X1 != 1){
+      run_dconv(dimA,
+                padA,
+                convstride1X1,
+                dilationA,
+                filterdimA1,
+                outdimA1,
+                CUDNN_DATA_HALF,
+                dx,
+                w,
+                dy1,
+                CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+      // add 2 together
+      grad_x.add_(grad_x_conv4);
+    }
+    else {
+      run_dconv_add(dimA,
+                    padA,
+                    convstride1X1,
+                    dilationA,
+                    filterdimA1,
+                    outdimA1,
+                    CUDNN_DATA_HALF,
+                    dx,
+                    w,
+                    dy1,
+                    dx_conv4);
+    }
+  }
+
+  DEBUG_MSG("[DEBUG] new dx : " << grad_x.to(at::kFloat).sum().item<float>());
+  DEBUG_MSG("[DEBUG] new wgrad1 : " << wgrad1.to(at::kFloat).sum().item<float>());
+  DEBUG_MSG("[DEBUG] new wgrad2 : " << wgrad2.to(at::kFloat).sum().item<float>());
+  DEBUG_MSG("[DEBUG] new wgrad3 : " << wgrad3.to(at::kFloat).sum().item<float>());
+  outputs.push_back(grad_x);
+  outputs.push_back(wgrad1);
+  outputs.push_back(wgrad2);
+  outputs.push_back(wgrad3);
+
+  if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+    DEBUG_MSG("[DEBUG] new wgrad4 : " << wgrad4.to(at::kFloat).sum().item<float>());
+    outputs.push_back(wgrad4);
+  }
+
+  return outputs;
+}
+
+namespace {
+
+struct bottleneck_forward_status {
+
+  int64_t dimA[4];
+  int64_t filterdimA1[4];
+  int64_t filterdimA2[4];
+  int64_t filterdimA3[4];
+  int64_t filterdimA4[4];
+
+  int axis[4];
+
+  int64_t outdimA0[4];
+  int64_t outdimA1[4];
+  int64_t outdimA2[4];
+  int64_t outdimA3[4];
+  int64_t outdimA4[4];
+
+  int64_t padA[2];
+  int64_t padA1[2];
+  int64_t padA2[2];  // halo padding
+  int64_t dilationA[2];
+  int64_t convstrideA[2];
+  int64_t convstride1X1[2];
+
+  int64_t outdim0[4]; // halo input shape
+  int64_t outdim1[4];
+  int64_t outdim2[4];
+  int64_t outdim3[4];
+  int64_t outdim4[4]; // halo output shape
+
+  void init(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs) {
+    dimA[0] = dimA[1] = dimA[2] = dimA[3] = 0;
+    filterdimA1[0] = filterdimA1[1] = filterdimA1[2] = filterdimA1[3] = 0;
+    filterdimA2[0] = filterdimA2[1] = filterdimA2[2] = filterdimA2[3] = 0;
+    filterdimA3[0] = filterdimA3[1] = filterdimA3[2] = filterdimA3[3] = 0;
+    filterdimA4[0] = filterdimA4[1] = filterdimA4[2] = filterdimA4[3] = 0;
+
+    // All dim calculation after this order of n,c,h,w
+    if (explicit_nhwc) {
+      axis[0] = 0;
+      axis[1] = 3;
+      axis[2] = 1;
+      axis[3] = 2;
+    } else {
+      axis[0] = 0;
+      axis[1] = 1;
+      axis[2] = 2;
+      axis[3] = 3;
+    }
+
+    for (int dim=0;dim<4;dim++) {
+      dimA[dim] = inputs[0].size(axis[dim]);
+      filterdimA1[dim] = inputs[1].size(axis[dim]);
+      filterdimA2[dim] = inputs[2].size(axis[dim]);
+      filterdimA3[dim] = inputs[3].size(axis[dim]);
+    }
+    if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+      for (int dim=0;dim<4;dim++) {
+	filterdimA4[dim] = inputs[10].size(axis[dim]);
+      }
+    }
+
+    // output dim in n,c,h,w used by backend
+    outdimA0[0] = outdimA0[1] = outdimA0[2] = outdimA0[3] = 0;
+    outdimA1[0] = outdimA1[1] = outdimA1[2] = outdimA1[3] = 0;
+    outdimA2[0] = outdimA2[1] = outdimA2[2] = outdimA2[3] = 0;
+    outdimA3[0] = outdimA3[1] = outdimA3[2] = outdimA3[3] = 0;
+    outdimA4[0] = outdimA4[1] = outdimA4[2] = outdimA4[3] = 0;
+
+    // use these fixed value for test run
+    padA[0] = 0; padA[1] = 0;
+    padA1[0] = 1; padA1[1] = 1;
+    padA2[0] = 0; padA2[1] = 1;
+    dilationA[0] = 1; dilationA[1] = 1;
+    convstrideA[0] = 1; convstrideA[1] = 1;
+    convstride1X1[0] = stride_1X1; convstride1X1[1] = stride_1X1;
+
+    // compute output from pad/stride/dilation
+    outdimA1[0] = dimA[0];
+    outdimA1[1] = filterdimA1[0];
+    for (int dim = 0; dim < 2; dim++) {
+      outdimA1[dim + 2] = getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);
+    }
+
+    outdimA2[0] = outdimA1[0];
+    outdimA2[1] = filterdimA2[0];
+    for (int dim = 0; dim < 2; dim++) {
+      outdimA2[dim + 2] = getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);
+    }
+
+    for (int dim = 0; dim < 4; dim++) {
+      if (dim == 2) {
+	outdimA0[dim] = 3;
+	outdimA4[dim] = 1;
+      } else {
+        outdimA0[dim] = outdimA1[dim];
+	outdimA4[dim] = outdimA2[dim];
+      }
+    }
+
+    outdimA3[0] = outdimA2[0];
+    outdimA3[1] = filterdimA3[0];
+    for (int dim = 0; dim < 2; dim++) {
+      outdimA3[dim + 2] = getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);
+    }
+
+    // Create output tensor in the correct shape in pytorch's view
+    outdim1[0] = outdim1[1] = outdim1[2] = outdim1[3] = 0;
+    outdim2[0] = outdim2[1] = outdim2[2] = outdim2[3] = 0;
+    outdim3[0] = outdim3[1] = outdim3[2] = outdim3[3] = 0;
+    if (explicit_nhwc) {
+      axis[0] = 0;
+      axis[1] = 2;
+      axis[2] = 3;
+      axis[3] = 1;
+    }
+    for (int dim=0;dim<4;dim++) {
+      outdim0[dim] = outdimA0[axis[dim]];
+      outdim1[dim] = outdimA1[axis[dim]];
+      outdim2[dim] = outdimA2[axis[dim]];
+      outdim3[dim] = outdimA3[axis[dim]];
+      outdim4[dim] = outdimA4[axis[dim]];
+    }
+  }
+};
+
+bottleneck_forward_status forward_state;
+
+} // end of anonymous namespace
+
+std::vector<at::Tensor> bottleneck_forward_init(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs) {
+  // NB! Bottleneck_forward and bottleneck_backward are NOT thread safe method.
+  // NB! We use a global object to store state.
+  forward_state.init(explicit_nhwc, stride_1X1, inputs);
+
+  // create output vector
+  std::vector<at::Tensor> outputs;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  //printf("outdim1 = (%d,%d,%d,%d)\n",forward_state.outdim1[0],forward_state.outdim1[1],forward_state.outdim1[2],forward_state.outdim1[3]);
+  auto out1 = at::empty(forward_state.outdim1, inputs[0].type(), output_format);
+  auto out2 = at::empty(forward_state.outdim2, inputs[0].type(), output_format);
+  auto out3 = at::empty(forward_state.outdim3, inputs[0].type(), output_format);
+
+  outputs.push_back(out1);
+  outputs.push_back(out2);
+  outputs.push_back(out3);
+
+  return outputs;
+}
+
+// inputs contains x,w,z,b,(i)
+void bottleneck_forward_out1(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs, std::vector<at::Tensor> outputs) {
+
+  std::cout << std::fixed;
+
+  // run
+  at::Half* x = inputs[0].data_ptr<at::Half>();
+  at::Half* w = inputs[1].data_ptr<at::Half>();
+  at::Half* z = inputs[4].data_ptr<at::Half>();
+  at::Half* b = inputs[7].data_ptr<at::Half>();
+  auto out1 = outputs[0];
+  at::Half* y1 = out1.data_ptr<at::Half>();
+
+  run_conv_scale_bias_add_activation(forward_state.dimA,
+                                     forward_state.padA,
+                                     forward_state.convstride1X1,
+                                     forward_state.dilationA,
+                                     forward_state.filterdimA1,
+                                     forward_state.outdimA1,
+                                     CUDNN_DATA_HALF,
+                                     x,
+                                     w,
+                                     y1,
+                                     z,
+                                     b,
+                                     nullptr);
+
+  DEBUG_MSG("[DEBUG] new relu1 : " << out1.to(at::kFloat).sum().item<float>());
+}
+
+// computes halo (top or bottom) from fat halo input.
+// fat halo input is 3 pixels wide in H.
+at::Tensor bottleneck_forward_out2_halo(bool explicit_nhwc, at::Tensor fat_halo_y1, std::vector<at::Tensor> inputs) {
+
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // run
+  at::Half* w = inputs[2].data_ptr<at::Half>();
+  at::Half* z = inputs[5].data_ptr<at::Half>();
+  at::Half* b = inputs[8].data_ptr<at::Half>();
+  
+  at::Half* y1 = fat_halo_y1.data_ptr<at::Half>();
+
+  auto halo_y2 = at::empty(forward_state.outdim4, inputs[0].type(), output_format);
+  at::Half* y2 = halo_y2.data_ptr<at::Half>();
+
+  run_conv_scale_bias_add_activation(forward_state.outdimA0,
+                                     forward_state.padA2,
+                                     forward_state.convstrideA,
+                                     forward_state.dilationA,
+                                     forward_state.filterdimA2,
+                                     forward_state.outdimA4,
+                                     CUDNN_DATA_HALF,
+                                     y1,
+                                     w,
+                                     y2,
+                                     z,
+                                     b,
+                                     nullptr);
+
+  return halo_y2;
+}
+
+void bottleneck_forward_out2(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs, std::vector<at::Tensor> outputs) {
+
+  std::cout << std::fixed;
+
+  // from _out1 method
+  at::Half* x = inputs[0].data_ptr<at::Half>();
+  auto out1 = outputs[0];
+  at::Half* y1 = out1.data_ptr<at::Half>();
+
+  // run
+  at::Half* w = inputs[2].data_ptr<at::Half>();
+  at::Half* z = inputs[5].data_ptr<at::Half>();
+  at::Half* b = inputs[8].data_ptr<at::Half>();
+  auto out2 = outputs[1];
+  at::Half* y2 = out2.data_ptr<at::Half>();
+
+  //printf("forward_state.outdimA1 = {%d,%d,%d,%d}\n",forward_state.outdimA1[0],forward_state.outdimA1[1],forward_state.outdimA1[2],forward_state.outdimA1[3]);
+  //printf("forward_state.padA1 = {%d,%d}\n",forward_state.padA1[0],forward_state.padA1[1]);
+  //printf("forward_state.convstrideA = {%d,%d}\n",forward_state.convstrideA[0],forward_state.convstrideA[1]);
+  //printf("forward_state.dilationA = {%d,%d}\n",forward_state.dilationA[0],forward_state.dilationA[1]);
+  //printf("forward_state.filterdimA2 = {%d,%d,%d,%d}\n",forward_state.filterdimA2[0],forward_state.filterdimA2[1],forward_state.filterdimA2[2],forward_state.filterdimA2[3]);
+  //printf("forward_state.outdimA2 = {%d,%d,%d,%d}\n",forward_state.outdimA2[0],forward_state.outdimA2[1],forward_state.outdimA2[2],forward_state.outdimA2[3]);
+  run_conv_scale_bias_add_activation(forward_state.outdimA1,
+                                     forward_state.padA1,
+                                     forward_state.convstrideA,
+                                     forward_state.dilationA,
+                                     forward_state.filterdimA2,
+                                     forward_state.outdimA2,
+                                     CUDNN_DATA_HALF,
+                                     y1,
+                                     w,
+                                     y2,
+                                     z,
+                                     b,
+                                     nullptr);
+  DEBUG_MSG("[DEBUG] new relu2 : " << out2.to(at::kFloat).sum().item<float>());
+}
+
+void bottleneck_forward_rest(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs, std::vector<at::Tensor> outputs) {
+
+  std::cout << std::fixed;
+
+  // from _out1 method
+  at::Half* x = inputs[0].data_ptr<at::Half>();
+
+  // create output of conv3
+  auto out3 = outputs[2];
+  at::Half* y3 = out3.data_ptr<at::Half>();
+
+  // create output of conv4 that may exist
+  auto identity = at::empty_like(out3);
+  at::Half* yi = identity.data_ptr<at::Half>();
+
+  at::Half *w, *z, *b;
+
+  if (stride_1X1 != 1 || forward_state.filterdimA3[0] != forward_state.dimA[1]){
+
+    w = inputs[10].data_ptr<at::Half>();
+    z = inputs[11].data_ptr<at::Half>();
+    b = inputs[12].data_ptr<at::Half>();
+    run_conv_scale_bias(forward_state.dimA,
+                        forward_state.padA,
+                        forward_state.convstride1X1,
+                        forward_state.dilationA,
+                        forward_state.filterdimA4,
+                        forward_state.outdimA3,
+                        CUDNN_DATA_HALF,
+                        x,
+                        w,
+                        yi,
+                        z,
+                        b);
+    DEBUG_MSG("[DEBUG] new downsample : " << identity.to(at::kFloat).sum().item<float>());
+  }
+  else {
+    yi = x;
+  }
+
+  auto out2 = outputs[1];
+  at::Half* y2 = out2.data_ptr<at::Half>();
+
+  w = inputs[3].data_ptr<at::Half>();
+  z = inputs[6].data_ptr<at::Half>();
+  b = inputs[9].data_ptr<at::Half>();
+
+  run_conv_scale_bias_add_activation(forward_state.outdimA2,
+                                     forward_state.padA,
+                                     forward_state.convstrideA,
+                                     forward_state.dilationA,
+                                     forward_state.filterdimA3,
+                                     forward_state.outdimA3,
+                                     CUDNN_DATA_HALF,
+                                     y2,
+                                     w,
+                                     y3,
+                                     z,
+                                     b,
+                                     yi);
+  DEBUG_MSG("[DEBUG] new relu3 : " << out3.to(at::kFloat).sum().item<float>());
+}
+
+namespace {
+
+struct bottleneck_backward_state {
+
+  int64_t dimA[4];
+  int64_t filterdimA1[4];
+  int64_t filterdimA2[4];
+  int64_t filterdimA3[4];
+  int64_t filterdimA4[4];
+  int64_t filterdimA2hh[4]; // Cin,Cout,1,3
+
+  int axis[4];
+
+  int64_t outdimA1[4]; // grad_out1
+  int64_t outdimA2[4]; // grad_out2
+  int64_t outdimA3[4];
+  int64_t outdimA1h[4]; // output: grad_out1 halo (H=3)
+  int64_t outdimA2h[4]; // input : grad_out2 halo cells (H=3)
+  int64_t outdimA1hh[4]; // input: grad_out2 halo (H=1)
+  int64_t outdimA2hh[4]; // input: out1 halo (H=1)
+
+  int64_t padA[2];
+  int64_t padA1[2];
+  int64_t padA2[2];
+  int64_t dilationA[2];
+  int64_t convstrideA[2];
+  int64_t convstride1X1[2];
+
+  int64_t filterdim2hh[4]; // Cin,1,3,Cout
+
+  int64_t outdim1[4];
+  int64_t outdim2[4];
+  int64_t outdim3[4];
+  int64_t outdim1h[4];
+
+  void init(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs) {
+    // setup dimensions
+    dimA[0] = dimA[1] = dimA[2] = dimA[3] = 0;
+    filterdimA1[0] = filterdimA1[1] = filterdimA1[2] = filterdimA1[3] = 0;
+    filterdimA2[0] = filterdimA2[1] = filterdimA2[2] = filterdimA2[3] = 0;
+    filterdimA3[0] = filterdimA3[1] = filterdimA3[2] = filterdimA3[3] = 0;
+    filterdimA4[0] = filterdimA4[1] = filterdimA4[2] = filterdimA4[3] = 0;
+    filterdimA2hh[0] = filterdimA2hh[1] = filterdimA2hh[2] = filterdimA2hh[3] = 0;
+
+    // All dim calculation after this order of n,c,h,w
+    if (explicit_nhwc) {
+      axis[0] = 0;
+      axis[1] = 3;
+      axis[2] = 1;
+      axis[3] = 2;
+    } else {
+      axis[0] = 0;
+      axis[1] = 1;
+      axis[2] = 2;
+      axis[3] = 3;
+    }
+
+    for (int dim=0;dim<4;dim++) {
+      dimA[dim] = inputs[0].size(axis[dim]);
+      filterdimA1[dim] = inputs[1].size(axis[dim]);
+      filterdimA2[dim] = inputs[2].size(axis[dim]);
+      filterdimA3[dim] = inputs[3].size(axis[dim]);
+    }
+    if (stride_1X1 != 1 || filterdimA3[0] != dimA[1]) {
+      for (int dim=0;dim<4;dim++) {
+        filterdimA4[dim] = inputs[14].size(axis[dim]);
+      }
+    }
+
+    for (int dim=0;dim<4;dim++) {
+      if (dim == 2) {
+        filterdimA2hh[dim] = 1;
+      } else {
+        filterdimA2hh[dim] = filterdimA2[dim];
+      }
+    }
+
+    // output dim in n,c,h,w used by backend
+    outdimA1[0] = outdimA1[1] = outdimA1[2] = outdimA1[3] = 0;
+    outdimA2[0] = outdimA2[1] = outdimA2[2] = outdimA2[3] = 0;
+    outdimA3[0] = outdimA3[1] = outdimA3[2] = outdimA3[3] = 0;
+    outdimA1h[0] = outdimA1h[1] = outdimA1h[2] = outdimA1h[3] = 0;
+    outdimA2h[0] = outdimA2h[1] = outdimA2h[2] = outdimA2h[3] = 0;
+    outdimA1hh[0] = outdimA1hh[1] = outdimA1hh[2] = outdimA1hh[3] = 0;
+    outdimA2hh[0] = outdimA2hh[1] = outdimA2hh[2] = outdimA2hh[3] = 0;
+
+    // use these fixed value for test run
+    padA[0] = 0; padA[1] = 0;
+    padA1[0] = 1; padA1[1] = 1;
+    padA2[0] = 0; padA2[1] = 1;
+    dilationA[0] = 1; dilationA[1] = 1;
+    convstrideA[0] = 1; convstrideA[1] = 1;
+    convstride1X1[0] = stride_1X1; convstride1X1[1] = stride_1X1;
+
+    // compute output from pad/stride/dilation
+    outdimA1[0] = dimA[0];
+    outdimA1[1] = filterdimA1[0];
+    for (int dim = 0; dim < 2; dim++) {
+      outdimA1[dim + 2] = getFwdConvOutputDim(dimA[dim + 2], padA[dim], filterdimA1[dim + 2], convstride1X1[dim], dilationA[dim]);
+    }
+
+    outdimA2[0] = outdimA1[0];
+    outdimA2[1] = filterdimA2[0];
+    for (int dim = 0; dim < 2; dim++) {
+      outdimA2[dim + 2] = getFwdConvOutputDim(outdimA1[dim + 2], padA1[dim], filterdimA2[dim + 2], convstrideA[dim], dilationA[dim]);
+    }
+
+    outdimA3[0] = outdimA2[0];
+    outdimA3[1] = filterdimA3[0];
+    for (int dim = 0; dim < 2; dim++) {
+      outdimA3[dim + 2] = getFwdConvOutputDim(outdimA2[dim + 2], padA[dim], filterdimA3[dim + 2], convstrideA[dim], dilationA[dim]);
+    }
+
+    for (int dim = 0; dim < 4; dim++) {
+      if (dim == 2) {
+	outdimA1h[dim] = 3;
+	outdimA2h[dim] = 3;
+        outdimA1hh[dim] = 1;
+        outdimA2hh[dim] = 1;
+      } else {
+	outdimA1h[dim] = outdimA1[dim];
+	outdimA2h[dim] = outdimA2[dim];
+        outdimA1hh[dim] = outdimA1[dim];
+        outdimA2hh[dim] = outdimA2[dim];
+      }
+    }
+
+    // Create output tensor in the correct shape in pytorch's view
+    outdim1[0] = outdim1[1] = outdim1[2] = outdim1[3] = 0;
+    outdim2[0] = outdim2[1] = outdim2[2] = outdim2[3] = 0;
+    outdim3[0] = outdim3[1] = outdim3[2] = outdim3[3] = 0;
+    outdim1h[0] = outdim1h[1] = outdim1h[2] = outdim1h[3] = 0;
+    filterdim2hh[0] = filterdim2hh[1] = filterdim2hh[2] = filterdim2hh[3] = 0;
+    if (explicit_nhwc) {
+      axis[0] = 0;
+      axis[1] = 2;
+      axis[2] = 3;
+      axis[3] = 1;
+    }
+    for (int dim=0;dim<4;dim++) {
+      outdim1[dim] = outdimA1[axis[dim]];
+      outdim2[dim] = outdimA2[axis[dim]];
+      outdim3[dim] = outdimA3[axis[dim]];
+      outdim1h[dim] = outdimA1h[axis[dim]];
+      filterdim2hh[dim] = filterdimA2hh[axis[dim]];
+    }
+  }
+};
+
+bottleneck_backward_state backward_state;
+
+}
+
+std::vector<at::Tensor> bottleneck_backward_init(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs) {
+
+  std::cout << std::fixed;
+
+  backward_state.init(explicit_nhwc, stride_1X1, inputs);
+
+  // create output vector
+  std::vector<at::Tensor> outputs;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  auto grad_x = at::empty_like(inputs[0]);
+  auto wgrad1 = at::empty_like(inputs[1]);
+  auto wgrad2 = at::empty_like(inputs[2]);
+  auto wgrad3 = at::empty_like(inputs[3]);
+
+  outputs.push_back(grad_x);
+  outputs.push_back(wgrad1);
+  outputs.push_back(wgrad2);
+  outputs.push_back(wgrad3);
+  if (stride_1X1 != 1 || backward_state.filterdimA3[0] != backward_state.dimA[1]) {
+    auto wgrad4 = at::empty_like(inputs[14]);
+    outputs.push_back(wgrad4);
+  }
+
+  return outputs;
+}
+
+at::Tensor bottleneck_backward_grad_out2(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs, std::vector<at::Tensor> outputs) {
+
+  bool requires_grad = inputs[0].requires_grad();
+
+  std::cout << std::fixed;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // dconv3+drelu2+dscale2
+  at::Half* conv_in = inputs[13].data_ptr<at::Half>();
+  at::Half* dy3 = inputs[10].data_ptr<at::Half>();
+
+  DEBUG_MSG("[DEBUG] new dconv3 : " << inputs[10].to(at::kFloat).sum().item<float>());
+
+  // wgrad
+  auto wgrad3 = outputs[3];
+  at::Half* dw3 = wgrad3.data_ptr<at::Half>();
+  run_dconv(backward_state.outdimA2,
+            backward_state.padA,
+            backward_state.convstrideA,
+            backward_state.dilationA,
+            backward_state.filterdimA3,
+            backward_state.outdimA3,
+            CUDNN_DATA_HALF,
+            conv_in,
+            dw3,
+            dy3,
+            CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+  // dgrad
+  auto grad_out2 = at::empty(backward_state.outdim2, inputs[0].type(), output_format);
+  at::Half* dy2 = grad_out2.data_ptr<at::Half>();
+  at::Half* w = inputs[3].data_ptr<at::Half>();
+  at::Half* z = inputs[5].data_ptr<at::Half>();
+
+  at::Half* relu2 = inputs[13].data_ptr<at::Half>();
+
+  run_dconv_drelu_dscale(backward_state.outdimA2,
+                         backward_state.padA,
+                         backward_state.convstrideA,
+                         backward_state.dilationA,
+                         backward_state.filterdimA3,
+                         backward_state.outdimA3,
+                         CUDNN_DATA_HALF,
+                         dy2,
+                         w,
+                         dy3,
+                         z,
+                         relu2);
+
+  // do halo exchange of dy2 here
+
+  DEBUG_MSG("[DEBUG] new dconv2 : " << grad_out2.to(at::kFloat).sum().item<float>());
+
+  return grad_out2;
+}
+
+at::Tensor bottleneck_backward_grad_out1(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs, std::vector<at::Tensor> outputs, at::Tensor grad_out2) {
+
+  bool requires_grad = inputs[0].requires_grad();
+
+  std::cout << std::fixed;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // dgrad
+  at::Half* dy2 = grad_out2.data_ptr<at::Half>();
+
+  // dgrad
+  auto grad_out1 = at::empty(backward_state.outdim1, inputs[0].type(), output_format);
+  at::Half* dy1 = grad_out1.data_ptr<at::Half>();
+  at::Half* w = inputs[2].data_ptr<at::Half>();
+  at::Half* z = inputs[4].data_ptr<at::Half>();
+
+  at::Half* relu1 = inputs[12].data_ptr<at::Half>();
+  //printf("relu.shape = [%d,%d,%d,%d]\n",inputs[12].size(0),inputs[12].size(1),inputs[12].size(2),inputs[12].size(3));
+
+  // fused dgrad
+  run_dconv_drelu_dscale(backward_state.outdimA1,
+                         backward_state.padA1,
+                         backward_state.convstrideA,
+                         backward_state.dilationA,
+                         backward_state.filterdimA2,
+                         backward_state.outdimA2,
+                         CUDNN_DATA_HALF,
+                         dy1,
+                         w,
+                         dy2,
+                         z,
+                         relu1);
+
+  return grad_out1;
+}
+
+// perform backward data 3x3 convolution (grad_out * w_rot180) on grad_out2 input of shape [N,3,W,C] with padding=(1,1) to produce output of shape [N,3,W,C]
+at::Tensor bottleneck_backward_grad_out1_halo(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs, std::vector<at::Tensor> outputs, at::Tensor grad_out2_halo, at::Tensor relu1_halo) {
+
+  bool requires_grad = inputs[0].requires_grad();
+
+  std::cout << std::fixed;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // dgrad
+  at::Half* dy2h = grad_out2_halo.data_ptr<at::Half>();
+
+  // dgrad
+  auto grad_out1_halo = at::empty(backward_state.outdim1h, inputs[0].type(), output_format);
+  at::Half* dy1h = grad_out1_halo.data_ptr<at::Half>();
+  at::Half* w = inputs[2].data_ptr<at::Half>();
+  at::Half* z = inputs[4].data_ptr<at::Half>();
+
+  at::Half* relu1h = relu1_halo.data_ptr<at::Half>();
+  //printf("relu.shape = [%d,%d,%d,%d]\n",relu1_halo.size(0),relu1_halo.size(1),relu1_halo.size(2),relu1_halo.size(3));
+  // fused dgrad
+  //printf("backward_state.outdimA1h = {%d,%d,%d,%d}\n",backward_state.outdimA1h[0],backward_state.outdimA1h[1],backward_state.outdimA1h[2],backward_state.outdimA1h[3]);
+  //printf("backward_state.outdimA2h = {%d,%d,%d,%d}\n",backward_state.outdimA2h[0],backward_state.outdimA2h[1],backward_state.outdimA2h[2],backward_state.outdimA2h[3]);
+  //printf("backward_state.filterdimA2 = {%d,%d,%d,%d}\n",backward_state.filterdimA2[0],backward_state.filterdimA2[1],backward_state.filterdimA2[2],backward_state.filterdimA2[3]);
+  run_dconv_drelu_dscale(backward_state.outdimA1h,
+                         backward_state.padA1,
+                         backward_state.convstrideA,
+                         backward_state.dilationA,
+                         backward_state.filterdimA2,
+                         backward_state.outdimA2h,
+                         CUDNN_DATA_HALF,
+                         dy1h,
+                         w,
+                         dy2h,
+                         z,
+                         relu1h);
+
+  return grad_out1_halo;
+}
+
+at::Tensor bottleneck_backward_wgrad2(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs, std::vector<at::Tensor> outputs, at::Tensor grad_out2) {
+
+  bool requires_grad = inputs[0].requires_grad();
+
+  std::cout << std::fixed;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // dgrad
+  at::Half* dy2 = grad_out2.data_ptr<at::Half>();
+
+  // dconv2+drelu1+dscale1
+  at::Half* conv_in = inputs[12].data_ptr<at::Half>();
+
+  // wgrad
+  auto wgrad2 = outputs[2];
+  at::Half* dw2 = wgrad2.data_ptr<at::Half>();
+  
+  //printf("outdimA1 = (%d,%d,%d,%d)\n",backward_state.outdimA1[0],backward_state.outdimA1[1],backward_state.outdimA1[2],backward_state.outdimA1[3]);
+  run_dconv(backward_state.outdimA1,
+            backward_state.padA1,
+            backward_state.convstrideA,
+            backward_state.dilationA,
+            backward_state.filterdimA2,
+            backward_state.outdimA2,
+            CUDNN_DATA_HALF,
+            conv_in,
+            dw2,
+            dy2,
+            CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+  return wgrad2;
+}
+
+// compute halo cells for input volume of dimension [N,1,W,C] with padding=(0,1) to produce output volume of dimension [N,1,W,C]
+// input and grad_out2_halo tensors are all of same shape
+// output tensor is of shape [Cin,1,3,Cout] (regular filter dims are [Cin,3,3,Cout]
+at::Tensor bottleneck_backward_wgrad2_halo(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs, std::vector<at::Tensor> outputs, at::Tensor input, at::Tensor grad_out2_halo) {
+
+  bool requires_grad = inputs[0].requires_grad();
+
+  std::cout << std::fixed;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // dgrad
+  at::Half* dy2 = grad_out2_halo.data_ptr<at::Half>();
+
+  // dconv2+drelu1+dscale1
+  at::Half* conv_in = input.data_ptr<at::Half>();
+
+  // wgrad
+  auto wgrad2_halo = at::empty(backward_state.filterdim2hh, input.type(), output_format);
+  at::Half* dw2 = wgrad2_halo.data_ptr<at::Half>();
+
+  //printf("backward_state.outdimA1hh = {%d,%d,%d,%d}\n",backward_state.outdimA1hh[0],backward_state.outdimA1hh[1],backward_state.outdimA1hh[2],backward_state.outdimA1hh[3]);
+  //printf("backward_state.outdimA2hh = {%d,%d,%d,%d}\n",backward_state.outdimA2hh[0],backward_state.outdimA2hh[1],backward_state.outdimA2hh[2],backward_state.outdimA2hh[3]);
+  //printf("backward_state.filterdim2hh = {%d,%d,%d,%d}\n",backward_state.filterdim2hh[0],backward_state.filterdim2hh[1],backward_state.filterdim2hh[2],backward_state.filterdim2hh[3]);
+  //printf("backward_state.filterdimA2hh = {%d,%d,%d,%d}\n",backward_state.filterdimA2hh[0],backward_state.filterdimA2hh[1],backward_state.filterdimA2hh[2],backward_state.filterdimA2hh[3]);
+  //printf("backward_state.padA2 = {%d,%d}\n",backward_state.padA2[0],backward_state.padA2[1]);
+  run_dconv(backward_state.outdimA1hh,  // N,C,1,W
+            backward_state.padA2, // 0, 1
+            backward_state.convstrideA,
+            backward_state.dilationA,
+            backward_state.filterdimA2hh,  // Cin,Cout,1,3
+            backward_state.outdimA2hh,  // N,C,1,W
+            CUDNN_DATA_HALF,
+            conv_in,
+            dw2,
+            dy2,
+            CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+  return wgrad2_halo;
+}
+
+void bottleneck_backward_rest(bool explicit_nhwc, int stride_1X1, std::vector<at::Tensor> inputs, std::vector<at::Tensor> outputs, at::Tensor grad_out2, at::Tensor grad_out1, at::Tensor wgrad2) {
+
+  bool requires_grad = inputs[0].requires_grad();
+
+  std::cout << std::fixed;
+  auto output_format = explicit_nhwc ? at::MemoryFormat::Contiguous : at::MemoryFormat::ChannelsLast;
+
+  // dgrad
+  at::Half* dy2 = grad_out2.data_ptr<at::Half>();
+  at::Half* dy1 = grad_out1.data_ptr<at::Half>();
+
+/*
+  // backward strided conv cannot be fused
+  // if stride == 1 but channel changes, we can fuse here
+  if (stride_1X1 != 1){
+    // dgrad
+    run_dconv(outdimA1,
+              padA1,
+              convstride1X1,
+              dilationA,
+              filterdimA2,
+              outdimA2,
+              CUDNN_DATA_HALF,
+              dy1,
+              w,
+              dy2,
+              CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+
+    // mul fused mask
+    grad_out1.mul_(inputs[15]);
+  }
+  else {
+    at::Half* relu1 = inputs[12].data_ptr<at::Half>();
+    // fused dgrad
+    run_dconv_drelu_dscale(outdimA1,
+                           padA1,
+                           convstride1X1,
+                           dilationA,
+                           filterdimA2,
+                           outdimA2,
+                           CUDNN_DATA_HALF,
+                           dy1,
+                           w,
+                           dy2,
+                           z,
+                           relu1);
+  }
+*/
+  DEBUG_MSG("[DEBUG] new dconv1 : " << grad_out1.to(at::kFloat).sum().item<float>());
+
+  // create grads of conv4 that may exist
+  auto grad_x_conv4 = at::empty_like(inputs[0]);
+  at::Half* dx_conv4 = grad_x_conv4.data_ptr<at::Half>();
+  at::Tensor wgrad4;
+
+  // x used for dconv1 and dconv4 wgrad
+  at::Half* x = inputs[0].data_ptr<at::Half>();
+
+  at::Half* w = NULL;
+
+  if (stride_1X1 != 1 || backward_state.filterdimA3[0] != backward_state.dimA[1]){
+    w = inputs[14].data_ptr<at::Half>();
+    at::Half* dy_conv4 = inputs[11].data_ptr<at::Half>();
+    if (requires_grad) {
+      run_dconv(backward_state.dimA,
+                backward_state.padA,
+                backward_state.convstride1X1,
+                backward_state.dilationA,
+                backward_state.filterdimA4,
+                backward_state.outdimA3,
+                CUDNN_DATA_HALF,
+                dx_conv4,
+                w,
+                dy_conv4,
+                CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+      // we don't print here since we can't hook out this grad in pytorch alone to compare, due to addition with dx
+      // DEBUG_MSG("[DEBUG] new dx_identity : " << grad_x_conv4.to(at::kFloat).sum().item<float>());
+    }
+    // wgrad
+    wgrad4 = outputs[4];
+    at::Half* dw4 = wgrad4.data_ptr<at::Half>();
+    run_dconv(backward_state.dimA,
+              backward_state.padA,
+              backward_state.convstride1X1,
+              backward_state.dilationA,
+              backward_state.filterdimA4,
+              backward_state.outdimA3,
+              CUDNN_DATA_HALF,
+              x,
+              dw4,
+              dy_conv4,
+              CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+  }
+  else {
+    // if there is no downsample, dx_conv4 is fork of drelu3
+    dx_conv4 = inputs[11].data_ptr<at::Half>();
+  }
+
+  // dconv1+add
+  // wgrad
+  auto wgrad1 = outputs[1];
+  at::Half* dw1 = wgrad1.data_ptr<at::Half>();
+  run_dconv(backward_state.dimA,
+            backward_state.padA,
+            backward_state.convstride1X1,
+            backward_state.dilationA,
+            backward_state.filterdimA1,
+            backward_state.outdimA1,
+            CUDNN_DATA_HALF,
+            x,
+            dw1,
+            dy1,
+            CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_FILTER_DESCRIPTOR);
+
+  // dgrad
+  w = inputs[1].data_ptr<at::Half>();
+  auto grad_x = outputs[0];
+  at::Half* dx = grad_x.data_ptr<at::Half>();
+
+  // backward strided conv cannot be fused
+  // if stride == 1 but channel changes, we can fuse here
+  if (requires_grad){
+    if (stride_1X1 != 1){
+      run_dconv(backward_state.dimA,
+                backward_state.padA,
+                backward_state.convstride1X1,
+                backward_state.dilationA,
+                backward_state.filterdimA1,
+                backward_state.outdimA1,
+                CUDNN_DATA_HALF,
+                dx,
+                w,
+                dy1,
+                CUDNN_BACKEND_OPERATION_CONVOLUTION_BACKWARD_DATA_DESCRIPTOR);
+      // add 2 together
+      grad_x.add_(grad_x_conv4);
+    }
+    else {
+      run_dconv_add(backward_state.dimA,
+                    backward_state.padA,
+                    backward_state.convstride1X1,
+                    backward_state.dilationA,
+                    backward_state.filterdimA1,
+                    backward_state.outdimA1,
+                    CUDNN_DATA_HALF,
+                    dx,
+                    w,
+                    dy1,
+                    dx_conv4);
+    }
+  }
+
+  DEBUG_MSG("[DEBUG] new dx : " << grad_x.to(at::kFloat).sum().item<float>());
+  DEBUG_MSG("[DEBUG] new wgrad1 : " << wgrad1.to(at::kFloat).sum().item<float>());
+  DEBUG_MSG("[DEBUG] new wgrad2 : " << wgrad2.to(at::kFloat).sum().item<float>());
+  DEBUG_MSG("[DEBUG] new wgrad3 : " << wgrad3.to(at::kFloat).sum().item<float>());
+
+  if (stride_1X1 != 1 || backward_state.filterdimA3[0] != backward_state.dimA[1]) {
+    DEBUG_MSG("[DEBUG] new wgrad4 : " << wgrad4.to(at::kFloat).sum().item<float>());
+  }
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &bottleneck_forward, "Bottleneck block forward");
+  m.def("backward", &bottleneck_backward, "Bottleneck block backward");
+  m.def("forward_init", &bottleneck_forward_init, "Bottleneck block init");
+  m.def("forward_out1", &bottleneck_forward_out1, "Bottleneck block forward");
+  m.def("forward_out2", &bottleneck_forward_out2, "Bottleneck block forward");
+  m.def("forward_out2_halo", &bottleneck_forward_out2_halo, "Bottleneck block forward");
+  m.def("forward_rest", &bottleneck_forward_rest, "Bottleneck block forward");
+  m.def("backward_init", &bottleneck_backward_init, "Bottleneck block backward init");
+  m.def("backward_grad_out2", &bottleneck_backward_grad_out2, "Bottleneck block backward");
+  m.def("backward_grad_out1", &bottleneck_backward_grad_out1, "Bottleneck block backward");
+  m.def("backward_grad_out1_halo", &bottleneck_backward_grad_out1_halo, "Bottleneck block backward");
+  m.def("backward_wgrad2", &bottleneck_backward_wgrad2, "Bottleneck block backward");
+  m.def("backward_wgrad2_halo", &bottleneck_backward_wgrad2_halo, "Bottleneck block backward");
+  m.def("backward_rest", &bottleneck_backward_rest, "Bottleneck block backward");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/fmha_api.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/fmha_api.cpp
new file mode 100644
index 0000000000..5e0d7d90c0
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/fmha_api.cpp
@@ -0,0 +1,432 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include <torch/extension.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include "fmha.h"
+
+void set_params(Fused_multihead_attention_fprop_params &params,
+                // sizes
+                const size_t b,
+                const size_t s,
+                const size_t h,
+                const size_t d,
+                // device pointers
+                void *qkv_packed_d,
+                void *cu_seqlens_d,
+                void *o_packed_d,
+                void *s_d,
+                float p_dropout) {
+
+    Data_type acc_type = DATA_TYPE_FP32;
+    Data_type data_type = DATA_TYPE_FP16;
+
+    // Reset the parameters
+    memset(&params, 0, sizeof(params));
+
+    // Set the pointers and strides.
+    params.qkv_ptr = qkv_packed_d;
+    params.qkv_stride_in_bytes = get_size_in_bytes(h * 3 * d, data_type);
+    params.o_ptr = o_packed_d;
+    params.o_stride_in_bytes = get_size_in_bytes(h * d, data_type);
+
+    params.cu_seqlens = static_cast<int *>(cu_seqlens_d);
+
+    // S = softmax(P)
+    params.s_ptr = s_d;
+    params.s_stride_in_bytes = get_size_in_bytes(b * h * s, data_type);
+
+    // Set the dimensions.
+    params.b = b;
+    params.h = h;
+    params.s = s;
+    params.d = d;
+
+    // Set the different scale values.
+    const float scale_bmm1 = 1.f / sqrtf(d);
+    constexpr float scale_softmax = 1.f;
+    constexpr float scale_bmm2 = 1.f;
+
+    set_alpha(params.scale_bmm1, scale_bmm1, acc_type);
+    set_alpha(params.scale_softmax, scale_softmax, acc_type);
+    set_alpha(params.scale_bmm2, scale_bmm2, data_type);
+
+    // Set this to probability of keeping an element to simplify things.
+    params.p_dropout = 1.f - p_dropout;
+    params.rp_dropout = 1.f / params.p_dropout;
+    TORCH_CHECK(p_dropout < 1.f);
+    set_alpha(params.scale_dropout, params.rp_dropout, data_type);
+}
+
+std::vector<at::Tensor>
+mha_fwd(const at::Tensor &qkv,  // total x num_heads x 3 x head_size, total := \sum_{i=0}^{b} s_i
+        const at::Tensor &cu_seqlens,  // b+1
+        const float p_dropout,
+        const int max_seq_len,
+        const bool is_training,
+        c10::optional<at::Generator> gen_) {
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+    TORCH_CHECK(dprops->major == 8 && dprops->minor == 0);
+    int seq_len = 512;
+    auto launch = &run_fmha_fp16_512_64_sm80;
+    if( max_seq_len <= 128 ) {
+        seq_len = 128;
+        launch = &run_fmha_fp16_128_64_sm80;
+    } else if( max_seq_len <= 256 ) {
+        seq_len = 256;
+        launch = &run_fmha_fp16_256_64_sm80;
+    } else if( max_seq_len <= 384 ) {
+        seq_len = 384;
+        launch = &run_fmha_fp16_384_64_sm80;
+    } else if( max_seq_len <= 512 ) {
+        seq_len = 512;
+        launch = &run_fmha_fp16_512_64_sm80;
+    } else {
+        TORCH_CHECK(false);
+    }
+
+    constexpr int warps_m = 1;
+    constexpr int warps_n = 4;  // this leads to an upper bound
+    const int mmas_m = seq_len / 16 / warps_m;
+    const int mmas_n = seq_len / 16 / warps_n;
+    
+    const int elts_per_thread = 8 * mmas_m * mmas_n;
+
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    TORCH_CHECK(qkv.dtype() == torch::kFloat16);
+    TORCH_CHECK(cu_seqlens.dtype() == torch::kInt32);
+
+    TORCH_CHECK(qkv.is_cuda())
+    TORCH_CHECK(cu_seqlens.is_cuda())
+
+    TORCH_CHECK(qkv.is_contiguous())
+    TORCH_CHECK(cu_seqlens.is_contiguous())
+
+    TORCH_CHECK(cu_seqlens.dim() == 1);
+    TORCH_CHECK(qkv.dim() == 4);
+
+    const auto sizes = qkv.sizes();
+
+    TORCH_CHECK(sizes[THREE_DIM] == 3);
+
+    const int batch_size = cu_seqlens.numel() - 1;
+    const int total = sizes[TOTAL_DIM];
+    const int num_heads = sizes[H_DIM];
+    const int head_size = sizes[D_DIM];
+    TORCH_CHECK(batch_size > 0);
+    TORCH_CHECK(head_size == 64);
+    auto opts = qkv.options();
+
+    auto ctx = torch::empty({ total, num_heads, head_size }, opts);
+
+    auto s = torch::empty({ batch_size, num_heads, seq_len, seq_len }, opts);
+
+    auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(
+        gen_, at::cuda::detail::getDefaultCUDAGenerator());
+
+    Fused_multihead_attention_fprop_params params;
+
+    set_params(params,
+               batch_size,
+               seq_len,
+               num_heads,
+               head_size,
+               qkv.data_ptr(),
+               cu_seqlens.data_ptr(),
+               ctx.data_ptr(),
+               s.data_ptr(),
+               p_dropout);
+
+    // number of times random will be generated per thread, to offset philox counter in the random
+    // state
+    int64_t counter_offset = elts_per_thread;
+    at::PhiloxCudaState rng_engine_inputs;
+
+    if( is_training ) {
+        // See Note [Acquire lock when using random generators]
+        std::lock_guard<std::mutex> lock(gen->mutex_);
+        params.philox_args = gen->philox_cuda_state(counter_offset);
+    }
+
+    launch(params, is_training, stream);
+
+    return { ctx, s };
+}
+
+std::vector<at::Tensor>
+mha_bwd(const at::Tensor &dout,  // total x num_heads, x head_size
+        const at::Tensor &qkv,   // total x num_heads x 3 x head_size, total := \sum_{i=0}^{b} s_i
+        at::Tensor &softmax,     // b x h x s x s softmax and dmask - will be overwritten with dP
+        const at::Tensor &cu_seqlens,  // b+1
+        const float p_dropout,         // probability to drop
+        const int max_seq_len          // max sequence length to choose the kernel
+) {
+    auto dprops = at::cuda::getCurrentDeviceProperties();
+    TORCH_CHECK(dprops->major == 8 && dprops->minor == 0);
+    int seq_len = 512;
+    auto launch = &run_fmha_dgrad_fp16_512_64_sm80;
+    if( max_seq_len <= 128 ) {
+        seq_len = 128;
+        launch = &run_fmha_dgrad_fp16_128_64_sm80;
+    } else if( max_seq_len <= 256 ) {
+        seq_len = 256;
+        launch = &run_fmha_dgrad_fp16_256_64_sm80;
+    } else if( max_seq_len <= 384 ) {
+        seq_len = 384;
+        launch = &run_fmha_dgrad_fp16_384_64_sm80;
+    } else if( max_seq_len <= 512 ) {
+        seq_len = 512;
+        launch = &run_fmha_dgrad_fp16_512_64_sm80;
+    } else {
+        TORCH_CHECK(false);
+    }
+
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    TORCH_CHECK(qkv.dtype() == torch::kFloat16);
+    TORCH_CHECK(dout.dtype() == torch::kFloat16);
+    TORCH_CHECK(softmax.dtype() == torch::kFloat16);
+    TORCH_CHECK(cu_seqlens.dtype() == torch::kInt32);
+
+    TORCH_CHECK(qkv.is_cuda());
+    TORCH_CHECK(cu_seqlens.is_cuda());
+
+    TORCH_CHECK(qkv.is_contiguous());
+    TORCH_CHECK(cu_seqlens.is_contiguous());
+
+    TORCH_CHECK(cu_seqlens.dim() == 1);
+    TORCH_CHECK(qkv.dim() == 4);
+
+    const auto sizes = qkv.sizes();
+
+    TORCH_CHECK(sizes[THREE_DIM] == 3);
+
+    const int batch_size = cu_seqlens.numel() - 1;
+    const int num_heads = sizes[H_DIM];
+    const int head_size = sizes[D_DIM];
+    TORCH_CHECK(batch_size > 0);
+    TORCH_CHECK(head_size == 64);
+
+    auto dqkv = torch::empty_like(qkv);
+
+    Fused_multihead_attention_fprop_params params;
+
+    set_params(params,
+               batch_size,
+               seq_len,
+               num_heads,
+               head_size,
+               qkv.data_ptr(),
+               cu_seqlens.data_ptr(),
+               dout.data_ptr(),     // we set o_ptr to dout
+               softmax.data_ptr(),  // softmax gets overwritten by dP!
+               p_dropout);
+
+    // we're re-using these scales
+    Data_type acc_type = DATA_TYPE_FP32;
+    set_alpha(params.scale_bmm1, 1.f, acc_type);
+    set_alpha(params.scale_softmax, 1.f / sqrtf(head_size), acc_type);
+    set_alpha(params.scale_bmm2, 1.f, DATA_TYPE_FP16);
+    params.dqkv_ptr = dqkv.data_ptr();
+
+    launch(params, stream);
+    return { dqkv, softmax };
+}
+
+std::vector<at::Tensor> mha_fwd_nl(const at::Tensor &qkv,         // total x num_heads x 3 x head_size, total := \sum_{i=0}^{b} s_i
+                                const at::Tensor &cu_seqlens,  // b+1
+                                const float p_dropout,
+                                const int max_seq_len,
+                                const bool is_training,
+                                c10::optional<at::Generator> gen_) {
+    int seq_len = 512;
+    auto launch = &run_fmha_fp16_512_64_sm80_nl;
+    TORCH_CHECK(max_seq_len == seq_len);
+
+    constexpr int warps_m = 1;
+    constexpr int warps_n = 4;  // this leads to an upper bound
+    const int mmas_m = seq_len / 16 / warps_m;
+    const int mmas_n = seq_len / 16 / warps_n;
+    // static_assert( mmas_m == 32 );
+    // static_assert( mmas_n == 4 );
+    const int elts_per_thread = 8 * mmas_m * mmas_n;
+
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    TORCH_CHECK(qkv.is_cuda())
+    TORCH_CHECK(cu_seqlens.is_cuda())
+
+    TORCH_CHECK(qkv.is_contiguous())
+    TORCH_CHECK(cu_seqlens.is_contiguous())
+
+    TORCH_CHECK(cu_seqlens.dim() == 1);
+    TORCH_CHECK(qkv.dim() == 4);
+
+    const auto sizes = qkv.sizes();
+
+    TORCH_CHECK(sizes[THREE_DIM] == 3);
+
+    const int batch_size = cu_seqlens.numel() - 1;
+    const int total = sizes[TOTAL_DIM];
+    const int num_heads = sizes[H_DIM];
+    const int head_size = sizes[D_DIM];
+    TORCH_CHECK(batch_size > 0);
+    TORCH_CHECK(head_size == 64);
+    auto opts = qkv.options();
+
+    auto ctx = torch::empty({ total, num_heads, head_size }, opts);
+
+    auto s = torch::empty({ batch_size, num_heads, seq_len, seq_len }, opts);
+
+    auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(gen_, at::cuda::detail::getDefaultCUDAGenerator());
+
+    Fused_multihead_attention_fprop_params params;
+
+    set_params(params,
+               batch_size,
+               seq_len,
+               num_heads,
+               head_size,
+               qkv.data_ptr(),
+               cu_seqlens.data_ptr(),
+               ctx.data_ptr(),
+               s.data_ptr(),
+               p_dropout);
+
+    // number of times random will be generated per thread, to offset philox counter in the random
+    // state
+    int64_t counter_offset = elts_per_thread;
+    at::PhiloxCudaState rng_engine_inputs;
+
+    if( is_training ) {
+        // See Note [Acquire lock when using random generators]
+        std::lock_guard<std::mutex> lock(gen->mutex_);
+        params.philox_args = gen->philox_cuda_state(counter_offset);
+    }
+    int num_chunks = 3;
+    if(batch_size == 3) {
+        num_chunks = 2;
+    }
+
+    launch(params, is_training, num_chunks, stream);
+
+    return { ctx, s };
+}
+
+std::vector<at::Tensor> mha_bwd_nl(const at::Tensor &dout,        // total x num_heads, x head_size
+                                const at::Tensor &qkv,         // total x num_heads x 3 x head_size, total := \sum_{i=0}^{b} s_i
+                                at::Tensor &softmax,           // b x h x s x s softmax and dmask - will be overwritten with dP
+                                const at::Tensor &cu_seqlens,  // b+1
+                                const float p_dropout,         // probability to drop
+                                const int max_seq_len          // max sequence length to choose the kernel
+) {
+
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    TORCH_CHECK(qkv.is_cuda())
+    TORCH_CHECK(cu_seqlens.is_cuda())
+
+    TORCH_CHECK(qkv.is_contiguous())
+    TORCH_CHECK(cu_seqlens.is_contiguous())
+
+    TORCH_CHECK(cu_seqlens.dim() == 1);
+
+    TORCH_CHECK(qkv.dim() == 4);
+
+    const auto sizes = qkv.sizes();
+
+    TORCH_CHECK(sizes[THREE_DIM] == 3);
+
+    const int batch_size = cu_seqlens.numel() - 1;
+    
+    const int total = sizes[TOTAL_DIM];
+    const int num_heads = sizes[H_DIM];
+    const int head_size = sizes[D_DIM];
+    TORCH_CHECK(batch_size > 0);
+    TORCH_CHECK(head_size == 64);
+
+    int seq_len = 512;
+    auto launch = &run_fmha_dgrad_fp16_512_64_sm80_nl;
+
+    auto opts = qkv.options();
+
+    auto dqkv = torch::empty_like(qkv);
+
+    int num_chunks = 2;
+    if( batch_size == 1 ) {
+        num_chunks = 4;
+    }else if( batch_size == 2 ) {
+        num_chunks = 3;
+    }
+    auto dkv = torch::empty({total, num_chunks, 2, num_heads, head_size}, opts);
+
+    Fused_multihead_attention_fprop_params params;
+
+    set_params(params,
+               batch_size,
+               seq_len,
+               num_heads,
+               head_size,
+               qkv.data_ptr(),
+               cu_seqlens.data_ptr(),
+               dout.data_ptr(),     // o_ptr = dout
+               softmax.data_ptr(),  // softmax gets overwritten by dP!
+               p_dropout);
+
+    params.dkv_ptr = dkv.data_ptr();
+
+    Data_type acc_type = DATA_TYPE_FP32;
+    set_alpha(params.scale_bmm1, 1.f, acc_type);
+    set_alpha(params.scale_softmax, 1.f / sqrtf(head_size), acc_type);
+    set_alpha(params.scale_bmm2, 1.f, DATA_TYPE_FP16);
+    params.dqkv_ptr = dqkv.data_ptr();
+
+    launch(params, num_chunks, stream);
+
+    //SPLIT-K reduction of num_chunks dK, dV parts
+
+    // The equivalent of the following Pytorch code:
+    // using namespace torch::indexing;
+    // at::Tensor view_out = dqkv.index({Slice(), Slice(1, None, None)});
+    // torch::sum_out(view_out, dkv, 1);
+
+    const int hidden_size = num_heads * head_size;
+    fmha_run_noloop_reduce(
+        dqkv.data_ptr(), dkv.data_ptr(), cu_seqlens.data_ptr<int>(), hidden_size, batch_size, total, num_chunks, stream);
+
+    return { dqkv, softmax, dkv };
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.doc() = "Fused Multi-head Self-attention for BERT";  
+    m.def("fwd", &mha_fwd, "Forward pass");
+    m.def("bwd", &mha_bwd, "Backward pass");
+    m.def("fwd_nl", &mha_fwd_nl, "Forward pass (small-batch)");
+    m.def("bwd_nl", &mha_bwd_nl, "Backward pass (small-batch)");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha.h
new file mode 100644
index 0000000000..a7b5e35584
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha.h
@@ -0,0 +1,125 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include <cuda.h>
+#include <vector>
+
+#include <ATen/CUDAGeneratorImpl.h>
+#include <ATen/cuda/CUDAGraphsUtils.cuh>
+
+#include <fmha_utils.h>
+
+
+constexpr int TOTAL_DIM = 0;
+constexpr int THREE_DIM = 1;
+constexpr int H_DIM = 2;
+constexpr int D_DIM = 3;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Qkv_params {
+    // The QKV matrices.
+    void *qkv_ptr;
+
+    // The stride between rows of the Q, K and V matrices.
+    size_t qkv_stride_in_bytes;
+
+    // The number of heads.
+    int h;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Fused_multihead_attention_fprop_params : public Qkv_params {
+
+    // The dQKV matrices.
+    void *dqkv_ptr;
+
+    // Temporary for dKV.
+    void *dkv_ptr;
+
+    // The O matrix (output).
+    void *o_ptr;
+
+    // The stride between rows of O.
+    int64_t o_stride_in_bytes;
+
+    // The pointer to the S matrix, overwritten by the dP matrix (bwd).
+    void *s_ptr;
+    // The stride between rows of the S matrix.
+    int64_t s_stride_in_bytes;
+
+    // The dimensions.
+    int b, s, d;
+
+    // The scaling factors for the kernel.
+    uint32_t scale_bmm1, scale_softmax, scale_bmm2;
+
+    // array of length b+1 holding starting offset of each sequence.
+    int *cu_seqlens;
+
+    // The dropout probability (probability of keeping an activation).
+    float p_dropout;
+
+    // Scale factor of 1 / (1 - p_dropout).
+    float rp_dropout;
+
+    // Scale factor of 1 / (1 - p_dropout), in half2.
+    uint32_t scale_dropout;
+
+    // Random state.
+    at::PhiloxCudaState philox_args;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+void run_fmha_fp16_128_64_sm80(const Fused_multihead_attention_fprop_params &params, bool is_training, cudaStream_t stream);
+void run_fmha_fp16_256_64_sm80(const Fused_multihead_attention_fprop_params &params, bool is_training, cudaStream_t stream);
+void run_fmha_fp16_384_64_sm80(const Fused_multihead_attention_fprop_params &params, bool is_training, cudaStream_t stream);
+void run_fmha_fp16_512_64_sm80(const Fused_multihead_attention_fprop_params &params, bool is_training, cudaStream_t stream);
+
+void run_fmha_dgrad_fp16_128_64_sm80(const Fused_multihead_attention_fprop_params &params, cudaStream_t stream);
+void run_fmha_dgrad_fp16_256_64_sm80(const Fused_multihead_attention_fprop_params &params, cudaStream_t stream);
+void run_fmha_dgrad_fp16_384_64_sm80(const Fused_multihead_attention_fprop_params &params, cudaStream_t stream);
+void run_fmha_dgrad_fp16_512_64_sm80(const Fused_multihead_attention_fprop_params &params, cudaStream_t stream);
+
+void run_fmha_fp16_512_64_sm80_nl(const Fused_multihead_attention_fprop_params &params, const bool is_training, const int num_chunks, cudaStream_t stream); 
+
+void run_fmha_dgrad_fp16_512_64_sm80_nl(const Fused_multihead_attention_fprop_params &params, const int num_chunks, cudaStream_t stream);
+
+void fmha_run_noloop_reduce(void *out,
+                            const void *in,
+                            const int *cu_seqlens,
+                            const int hidden_size,
+                            const int batch_size,
+                            const int total,
+                            const int num_chunks,
+                            cudaStream_t stream);
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/gemm.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/gemm.h
new file mode 100644
index 0000000000..b6c9f2c9c3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/gemm.h
@@ -0,0 +1,317 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include <fmha/utils.h>
+
+#define FMHA_DIV_UP(m, n) (((m) + (n)-1) / (n))
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Data_type_, int NUM_ELTS_, int BITS_PER_ELT_, int ALIGNMENT_ >
+struct Fragment_base_ {
+
+    // The data type.
+    using Data_type = Data_type_;
+    // default input type
+    using Input_type_ = Data_type_;
+    // Does it store the array of elements.
+    enum { HAS_ELTS = BITS_PER_ELT_ >= 8 };
+    // The number of elements.
+    enum { NUM_ELTS = NUM_ELTS_ };
+    // The size of element in bits.
+    enum { BITS_PER_ELT = BITS_PER_ELT_ };
+    // The size of byte of a single register.
+    enum { BYTES_PER_REG = 4 };
+    // The size in bits.
+    enum { BITS_PER_REG = BYTES_PER_REG * 8 };
+    // The number of registers needed to store the fragment.
+    enum { NUM_REGS = Div_up<NUM_ELTS * BITS_PER_ELT, BITS_PER_REG>::VALUE };
+    // The size in bytes (as returned by sizeof(Fragment_base<>).
+    enum { SIZE_IN_BYTES = NUM_REGS * BYTES_PER_REG };
+    // The alignment.
+    enum { ALIGNMENT = ALIGNMENT_ > 0 ? ALIGNMENT_ : Min<NUM_REGS * BYTES_PER_REG, 16>::VALUE };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The type of the elements.
+    typename Data_type_,
+    // The number of elements.
+    int NUM_ELTS_,
+    // The alignment if you want to force a value -- use 0 otherwise.
+    int ALIGNMENT_ = 0,
+    // The base class.
+    typename Base_ = Fragment_base_<Data_type_, NUM_ELTS_, 8 * sizeof(Data_type_), ALIGNMENT_>
+>
+struct alignas(static_cast<int>(Base_::ALIGNMENT)) Fragment : public Base_ {
+
+    // The size of a load/store.
+    enum { BYTES_PER_LOAD_STORE = Base_::NUM_REGS * sizeof(uint32_t) };
+
+    // Clear the fragment. Using PTX in that code seems to produce better SASS...
+    inline __device__ void clear() {
+        #pragma unroll
+        for( int ii = 0; ii < Base_::NUM_REGS; ++ii ) {
+            asm volatile("mov.u32 %0, 0; \n" : "=r"(this->reg(ii)) : );
+        }
+    }
+
+    // Immutable access to a register.
+    inline __device__ const uint32_t& reg(int ii) const {
+        return this->regs_[ii];
+    }
+
+    // Mutable access to a register.
+    inline __device__ uint32_t& reg(int ii) {
+        return this->regs_[ii];
+    }
+
+    uint32_t regs_[Base_::NUM_REGS];
+
+    // Immutable access to the elements.
+    inline __device__ const Data_type_& elt(int ii) const {
+        return reinterpret_cast<const Data_type_*>(&this->regs_[0])[ii];
+    }
+
+    // Mutable access to the elements.
+    inline __device__ Data_type_& elt(int ii) {
+        return reinterpret_cast<Data_type_*>(&this->regs_[0])[ii];
+    }
+
+    // Immutable access to the elements with a cast.
+    template< typename Cast_type >
+    inline __device__ const Cast_type& elt_as(int ii) const {
+        return reinterpret_cast<const Cast_type*>(&this->regs_[0])[ii];
+    }
+
+    // Mutable access to the elements.
+    template< typename Cast_type >
+    inline __device__ Cast_type& elt_as(int ii) {
+        return reinterpret_cast<Cast_type*>(&this->regs_[0])[ii];
+    }
+
+    // Add another fragment.
+    inline __device__ void add(const Fragment &other) {
+        #pragma unroll
+        for( int ii = 0; ii < NUM_ELTS_; ++ii ) {
+            this->elt(ii) += other.elt(ii);
+        }
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Layout >
+struct Fragment_a : public Fragment<uint16_t, 8> {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Layout >
+struct Fragment_b : public Fragment<uint16_t, 8> {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Fragment_accumulator : public Fragment<float, 8> {
+
+    // The base class.
+    using Base = Fragment<float, 8>;
+
+    // Add two fragments.
+    template< typename Other_fragment_ >
+    inline __device__ void add(const Other_fragment_ &other) {
+        for( int ii = 0; ii < Base::NUM_ELTS; ++ii ) {
+            this->elt(ii) = this->elt(ii) + other.elt(ii);
+        }
+    }
+
+    // Do the HMMA.
+    template< typename Layout_a, typename Layout_b >
+    inline __device__ void mma(const Fragment_a<Layout_a> &a,
+                               const Fragment_b<Layout_b> &b) {
+        asm volatile( \
+            "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 \n" \
+            "    {%0, %1, %2, %3}, \n" \
+            "    {%4, %5, %6, %7}, \n" \
+            "    {%8, %9}, \n" \
+            "    {%0, %1, %2, %3}; \n" \
+                    : "+f"(  elt(0)), "+f"(  elt(1)), "+f"(  elt(2)), "+f"(  elt(3))
+                    :  "r"(a.reg(0)),  "r"(a.reg(1)),  "r"(a.reg(2)),  "r"(a.reg(3))
+                    ,  "r"(b.reg(0)),  "r"(b.reg(1)));
+        asm volatile( \
+            "mma.sync.aligned.m16n8k16.row.col.f32.f16.f16.f32 \n" \
+            "    {%0, %1, %2, %3}, \n" \
+            "    {%4, %5, %6, %7}, \n" \
+            "    {%8, %9}, \n" \
+            "    {%0, %1, %2, %3}; \n" \
+                    : "+f"(  elt(4)), "+f"(  elt(5)), "+f"(  elt(6)), "+f"(  elt(7))
+                    :  "r"(a.reg(0)),  "r"(a.reg(1)),  "r"(a.reg(2)),  "r"(a.reg(3))
+                    ,  "r"(b.reg(2)),  "r"(b.reg(3)));
+    }
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Fragment, int M, int N >
+inline __device__ void clear(Fragment (&frag)[M][N]) {
+    #pragma unroll
+    for( int mi = 0; mi < M; ++mi ) {
+        #pragma unroll
+        for( int ni = 0; ni < N; ++ni ) {
+            frag[mi][ni].clear();
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Accumulator_type, int WARPS_K >
+struct Clear_accumulator {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int WARPS_K >
+struct Clear_accumulator<float, WARPS_K> {
+  template< typename Acc, int M, int N >
+  static inline __device__ void apply(Acc (&acc)[M][N], bool = false) {
+    fmha::clear(acc);
+  }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Acc, typename A, typename B, int M, int N>
+inline __device__ void gemm(Acc (&acc)[M][N], const A (&a)[M], const B (&b)[N]) {
+
+    #pragma unroll
+    for( int mi = 0; mi < M; ++mi ) {
+        #pragma unroll
+        for( int ni = 0; ni < N; ++ni ) {
+            acc[mi][ni].mma(a[mi], b[ni]);
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The number of rows in the CTA tile.
+    int M_,
+    // The number of cols in the CTA tile.
+    int N_,
+    // The number of elements in the the K dimension of the GEMM loop.
+    int K_,
+    // The number of rows of warps.
+    int WARPS_M_,
+    // The number of cols of warps.
+    int WARPS_N_,
+    // The number of warps in the K dimension of the GEMM loop.
+    int WARPS_K_>
+struct Cta_tile_ {
+
+    enum { M = M_, N = N_, K = K_ };
+    // The number of warps.
+    enum { WARPS_M = WARPS_M_, WARPS_N = WARPS_N_, WARPS_K = WARPS_K_ };
+    // The number of warps per CTA.
+    enum { WARPS_PER_CTA = WARPS_M * WARPS_N * WARPS_K };
+    // The number of threads per warp.
+    enum { THREADS_PER_WARP = 32 };
+    // The number of threads per CTA.
+    enum { THREADS_PER_CTA = WARPS_PER_CTA * THREADS_PER_WARP };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Cta_tile>
+struct Hmma_tile {
+    // The number of elements computed with a single warp-MMA.
+    enum { M_PER_MMA = 16, N_PER_MMA = 16, K_PER_MMA = 16 };
+
+    // The number of elements computed with a single CTA-MMA.
+    enum {
+        M_PER_MMA_PER_CTA = M_PER_MMA * Cta_tile::WARPS_M,
+        N_PER_MMA_PER_CTA = N_PER_MMA * Cta_tile::WARPS_N,
+        K_PER_MMA_PER_CTA = K_PER_MMA * Cta_tile::WARPS_K
+    };
+
+    // The number of MMAs needed to compute the GEMM.
+    enum {
+        MMAS_M = Div_up<Cta_tile::M, M_PER_MMA_PER_CTA>::VALUE,
+        MMAS_N = Div_up<Cta_tile::N, N_PER_MMA_PER_CTA>::VALUE,
+        MMAS_K = Div_up<Cta_tile::K, K_PER_MMA_PER_CTA>::VALUE,
+    };
+
+    // The number of elements computed per warp.
+    enum {
+        M_PER_WARP = MMAS_M * M_PER_MMA,
+        N_PER_WARP = MMAS_N * N_PER_MMA,
+        K_PER_WARP = MMAS_K * K_PER_MMA,
+    };
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+using A_type = uint16_t;
+using B_type = uint16_t;
+using C_type = uint16_t;
+using Accumulator_type = float;
+using Epilogue_type = float;
+
+constexpr int BITS_PER_ELEMENT_A = sizeof(A_type) * 8;
+constexpr int BITS_PER_ELEMENT_B = sizeof(B_type) * 8;
+constexpr int BITS_PER_ELEMENT_C = sizeof(C_type) * 8;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int M, int N, int K, int WARPS_M, int WARPS_N, int WARPS_K>
+using Cta_tile_extd = Cta_tile_<M, N, K, WARPS_M, WARPS_N, WARPS_K>;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Cta_tile_>
+using Cta_tile_with_k_with_padding = Cta_tile_extd<Cta_tile_::M,
+                                                   Cta_tile_::N,
+                                                   Next_power_of_two<Cta_tile_::K>::VALUE,
+                                                   Cta_tile_::WARPS_M,
+                                                   Cta_tile_::WARPS_N,
+                                                   Cta_tile_::WARPS_K>;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/gmem_tile.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/gmem_tile.h
new file mode 100644
index 0000000000..7548a8e107
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/gmem_tile.h
@@ -0,0 +1,428 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The dimensions of the tile computed by the CTA.
+    typename Cta_tile,
+    // The number of bits per element.
+    int BITS_PER_ELEMENT,
+    // The number of rows of Q, K or V loaded by this tile.
+    int ROWS,
+    // The number of columns.
+    int COLS,
+    // The number of matrics.
+    int NUM_MATS = 3
+>
+struct Gmem_tile_qkv {
+
+    // The size of each LDG.
+    enum { BYTES_PER_LDG = 16 };
+    // The size of a row in bytes.
+    enum { BYTES_PER_ROW = COLS * BITS_PER_ELEMENT / 8 };
+
+    // The number of threads to load a "row" of the matrix.
+    enum { THREADS_PER_ROW = BYTES_PER_ROW / BYTES_PER_LDG };
+
+    // The number of "rows" loaded per LDG.
+    enum { ROWS_PER_LDG = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW };
+    // The number of LDGs needed to load a chunk of the Q matrix.
+    enum { LDGS = fmha::Div_up<ROWS, ROWS_PER_LDG>::VALUE };
+
+    // Ctor.
+    template< typename Params, typename BInfo >
+    inline __device__ Gmem_tile_qkv(const Params &params, int qkv_offset, const BInfo &binfo, int tidx)
+        : params_qkv_stride_in_bytes_(params.qkv_stride_in_bytes)
+        , actual_seqlen(binfo.actual_seqlen)
+        , qkv_ptr_(reinterpret_cast<char *>(params.qkv_ptr)) {
+
+        // Compute the position in the sequence (within the CTA for the moment).
+        int row = tidx / THREADS_PER_ROW;
+        // Compute the position of the thread in the row.
+        int col = tidx % THREADS_PER_ROW;
+
+        // Store the row as we need it to disable the loads.
+        row_ = row;
+
+        // The row offset in the batched GEMM. For each seq element, we store QKV in that order.
+        int64_t row_offset = (int64_t)row * params.qkv_stride_in_bytes;
+        // Add the block index.
+        row_offset += (int64_t)((binfo.sum_s * NUM_MATS + qkv_offset) * binfo.h + binfo.bidh) * BYTES_PER_ROW;
+
+        // Assemble the final pointer.
+        qkv_ptr_ += row_offset + col * BYTES_PER_LDG;
+    }
+
+    // Store data to shared memory.
+    template< typename Smem_tile >
+    inline __device__ void commit(Smem_tile &smem_tile) {
+        smem_tile.store(fetch_);
+    }
+
+    // Load data from memory.
+    template< typename Smem_tile >
+    inline __device__ void load(Smem_tile &smem_tile) {
+        const void *ptrs[LDGS];
+        uint32_t preds[LDGS];
+        #pragma unroll
+        for( int ii = 0; ii < LDGS; ++ii ) {
+            ptrs[ii] = qkv_ptr_ + (int64_t)ii * ROWS_PER_LDG * params_qkv_stride_in_bytes_;
+            preds[ii] = ((row_ + ii * ROWS_PER_LDG) < min(ROWS, actual_seqlen));
+            fetch_[ii] = make_uint4(0, 0, 0, 0);
+        }
+
+        // not packing predicates removes restrictions (e.g. FP16 384, 4 warps)
+        Ldg_functor<uint4, LDGS> fct(fetch_, ptrs);
+        #pragma unroll
+        for( int ii = 0; ii < LDGS; ++ii ) {
+            fct.load(ii, preds[ii]);
+        }
+    }
+
+    // Store data to memory.
+    inline __device__ void store(const uint4 (&data)[LDGS]) {
+        #pragma unroll
+        for( int ii = 0; ii < LDGS; ++ii ) {
+            char *ptr = qkv_ptr_ + (int64_t)ii * ROWS_PER_LDG * params_qkv_stride_in_bytes_;
+            if( (row_ + ii * ROWS_PER_LDG) < min(ROWS, actual_seqlen) ) {
+                fmha::stg(ptr, data[ii]);
+            }
+        }
+    }
+
+    // Move the pointer to the next location.
+    inline __device__ void move() {
+        qkv_ptr_ += (int64_t)ROWS * params_qkv_stride_in_bytes_;
+        actual_seqlen -= ROWS;
+    }
+
+    // The stride between rows for the QKV matrice.
+    int64_t params_qkv_stride_in_bytes_;
+    // The pointer.
+    char *qkv_ptr_;
+    // The fetch registers.
+    uint4 fetch_[LDGS];
+    // Keep track of the row the thread is processing as we move the tile.
+    int row_;
+    // The length of the sequence loaded by that memory tile.
+    int actual_seqlen;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Cta_tile >
+struct Gmem_tile_o {
+
+    // The mma tile.
+    using Mma_tile = fmha::Hmma_tile<Cta_tile>;
+
+    // The size of each element.
+    enum { BYTES_PER_ELEMENT = 2 };
+    // The size of a row in bytes.
+    enum { BYTES_PER_ROW = Cta_tile::N * BYTES_PER_ELEMENT };
+
+    // The number of threads to store a "row" of the matrix.
+    enum { THREADS_PER_ROW = 16 };
+    // The size of each STG.
+    enum { BYTES_PER_STG = BYTES_PER_ROW / THREADS_PER_ROW };
+
+    // The number of "rows" stored per iteration of the loop. The output of 1 MMA.
+    enum { ROWS = Cta_tile::M };
+    // The number of "rows" stored per iteration of the loop. The output of 1 MMA.
+    enum { ROWS_PER_LOOP = ROWS <= 64 ? ROWS : (int)Mma_tile::M_PER_MMA_PER_CTA };
+    // The number of outter loop for the stores.
+    enum { LOOPS = ROWS / ROWS_PER_LOOP };
+
+    // The number of "rows" stored per STG.
+    enum { ROWS_PER_STG = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW };
+    // Do we have to guard against partial writes/reads.
+    enum { HAS_INCOMPLETE_STG = Cta_tile::M % ROWS_PER_STG != 0 };
+    // The number of STGs needed to store a chunk of the Q matrix.
+    enum { STGS_PER_LOOP = fmha::Div_up<ROWS_PER_LOOP, ROWS_PER_STG>::VALUE };
+    // The number of STGs needed to store a chunk of the Q matrix in total.
+    enum { STGS = STGS_PER_LOOP * LOOPS };
+
+    // Ctor.
+    template<typename Params, typename BInfo>
+    inline __device__ Gmem_tile_o(const Params &params, const BInfo &binfo, int tidx)
+        : params_o_stride_in_bytes_(params.o_stride_in_bytes)
+        , actual_seqlen_(binfo.actual_seqlen)
+        , o_ptr_(reinterpret_cast<char *>(params.o_ptr)) {
+
+        // Compute the position in the sequence (within the CTA for the moment).
+        int row = tidx / THREADS_PER_ROW;
+        // Compute the position of the thread in the row.
+        int col = tidx % THREADS_PER_ROW;
+
+        // Store the row as we need it to disable loads.
+        row_ = row;
+
+        // The row offset in the batched GEMM.
+        int64_t row_offset = (int64_t)row * params.o_stride_in_bytes + binfo.bidx * BYTES_PER_ROW;
+        // Assemble the final pointer.
+        o_ptr_ += row_offset + col * BYTES_PER_STG;
+
+        // Is that thread active on the last STG?
+        if( HAS_INCOMPLETE_STG ) {
+            is_active_for_last_stg_ = row + (STGS - 1) * ROWS_PER_STG < Cta_tile::M;
+        }
+    }
+
+    // Store data to global memory.
+    inline __device__ void store(const uint4 (&src)[STGS_PER_LOOP], int mi) {
+
+        #pragma unroll
+        for( int ii = 0; ii < STGS_PER_LOOP; ++ii ) {
+            int jj = mi * STGS_PER_LOOP + ii;
+            if( this->row_ + jj * ROWS_PER_STG >= this->actual_seqlen_ ) {
+                break;
+            }
+
+            float x = reinterpret_cast<const float &>(src[ii].x);
+            float y = reinterpret_cast<const float &>(src[ii].y);
+            float z = reinterpret_cast<const float &>(src[ii].z);
+            float w = reinterpret_cast<const float &>(src[ii].w);
+            uint2 out = float4_to_half4(x, y, z, w);
+            if( !HAS_INCOMPLETE_STG || (jj < STGS - 1 || this->is_active_for_last_stg_) ) {
+                fmha::stg(this->o_ptr_ + jj * ROWS_PER_STG * this->params_o_stride_in_bytes_, out);
+            }
+        }
+    }
+
+    // Move the pointer to the next location.
+    inline __device__ void move() {
+        row_ += ROWS;
+        o_ptr_ += (int64_t)ROWS * params_o_stride_in_bytes_;
+    }
+
+    // The stride between rows for the QKV matrice.
+    int64_t params_o_stride_in_bytes_;
+    // The pointer.
+    char *o_ptr_;
+    // Is the thread active for the last STG?
+    int is_active_for_last_stg_;
+    // Keep track of the row to disable loads.
+    int row_;
+    // The length of the sequence loaded by that memory tile.
+    int actual_seqlen_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Cta_tile, int BYTES_PER_ELEMENT >
+struct Gmem_tile_mma_sd {
+
+    // The mma tile.
+    using Mma_tile = fmha::Hmma_tile<Cta_tile>;
+
+    // Each STG stores 8 elements.
+    enum { BYTES_PER_STG = BYTES_PER_ELEMENT * 8 };
+    // The number of MMAs in the M dimension.
+    enum { MMAS_M = Mma_tile::MMAS_M };
+    // The number of MMAs in the N dimension.
+    enum { MMAS_N = Mma_tile::MMAS_N };
+    // The number of rows computed per MMA per thread block.
+    enum { M_PER_MMA_PER_CTA = Mma_tile::M_PER_MMA_PER_CTA };
+    // The number of cols computed per MMA per thread block.
+    enum { N_PER_MMA_PER_CTA = Mma_tile::N_PER_MMA_PER_CTA };
+    // The number of threads per block.
+    enum { THREADS_PER_CTA = Cta_tile::THREADS_PER_CTA };
+    // The size of each row in bytes. I.e. how many bytes are stored per STG.
+    enum { BYTES_PER_ROW = THREADS_PER_CTA * BYTES_PER_STG };
+    // The fixed sequence length.
+    enum { SEQLEN = Cta_tile::N };
+    // The distance between two blocks (in bytes).
+    enum { BLOCK_STRIDE_BYTES = SEQLEN * SEQLEN * BYTES_PER_ELEMENT };
+    // The distance between elements stored per loop (in bytes).
+    enum { LOOP_STRIDE_BYTES = MMAS_M * MMAS_N * BYTES_PER_ROW };
+
+    // The type of elements stored per STG.
+    using Type = typename fmha::Uint_from_size_in_bytes<BYTES_PER_STG>::Type;
+
+    // Ctor.
+    template<typename Params>
+    inline __device__ Gmem_tile_mma_sd(void *ptr, const Params &params, const int tidx) 
+        : ptr_(static_cast<char *>(ptr)) {
+
+        // The block index for the batch.
+        const int bidb = blockIdx.y;
+        // The block index for the head.
+        const int bidh = blockIdx.x;
+        // The block index.
+        size_t bidx = bidb * params.h + bidh;
+
+        // Set store location for each thread at the beginning of the loop
+        ptr_ += bidx * BLOCK_STRIDE_BYTES + tidx * BYTES_PER_STG;
+    }
+
+    // Store to global memory.
+    inline __device__ void store(const Type &data, const int mi, const int ni) {
+        size_t offset = (mi * MMAS_N + ni) * BYTES_PER_ROW;
+        fmha::stg(ptr_ + offset, data);
+    }
+
+    // Load from global memory.
+    inline __device__ void load(Type &data, const int mi, const int ni) {
+        size_t offset = (mi * MMAS_N + ni) * BYTES_PER_ROW;
+        fmha::ldg(data, ptr_ + offset);
+    }
+
+    // Move to the next tile.
+    inline __device__ void move() {
+        ptr_ += LOOP_STRIDE_BYTES;
+    }
+
+    // The pointer in global memory.
+    char *ptr_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Cta_tile, typename Base = Gmem_tile_mma_sd<Cta_tile, sizeof(uint16_t)> >
+struct Gmem_tile_mma_s : public Base {
+
+    // The number of mmas in the vertical dimension.
+    enum { M = Base::MMAS_M };
+    // The number of mmas in the horizontal dimension.
+    enum { N = Base::MMAS_N };
+    // The type of the vectors stored by each STG.
+    using Type = typename Base::Type;
+
+    // Ctor.
+    template< typename Params >
+    inline __device__ Gmem_tile_mma_s(void *ptr, const Params &params, const int tidx) 
+        : Base(ptr, params, tidx) {
+    }
+
+    // Store to global memory.
+    template<typename Mask>
+    inline __device__ void store(const float (&softmax)[2 * M][4 * N], const Mask &mask) {
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ni = 0; ni < N; ni++ ) {
+
+                float tmp00 = softmax[2 * mi + 0][4 * ni + 0];
+                float tmp01 = softmax[2 * mi + 0][4 * ni + 1];
+                float tmp02 = softmax[2 * mi + 0][4 * ni + 2];
+                float tmp03 = softmax[2 * mi + 0][4 * ni + 3];
+
+                float tmp10 = softmax[2 * mi + 1][4 * ni + 0];
+                float tmp11 = softmax[2 * mi + 1][4 * ni + 1];
+                float tmp12 = softmax[2 * mi + 1][4 * ni + 2];
+                float tmp13 = softmax[2 * mi + 1][4 * ni + 3];
+
+                uint4 dst;
+                dst.x = fmha::float2_to_half2(tmp00, tmp01);
+                dst.y = fmha::float2_to_half2(tmp02, tmp03);
+                dst.z = fmha::float2_to_half2(tmp10, tmp11);
+                dst.w = fmha::float2_to_half2(tmp12, tmp13);
+                if( mask.is_valid(mi, ni, 0, 0) ) {
+                    Base::store(dst, mi, ni);
+                }
+            }
+        }
+    }
+
+    // Load from global memory.
+    template<typename Mask>
+    inline __device__ void load(uint4 (&regs)[M][N], const Mask &mask) {
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ni = 0; ni < N; ni++ ) {
+                regs[mi][ni] = make_uint4(0, 0, 0, 0);
+                if( mask.is_valid(mi, ni, 0, 0) ) {
+                    Base::load(regs[mi][ni], mi, ni);
+                }
+            }
+        }
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The dimensions of the tile computed by the CTA.
+    typename Cta_tile,
+    // The base class.
+    typename Base = fmha::Gmem_tile_qkv<Cta_tile, fmha::BITS_PER_ELEMENT_A, Cta_tile::M, Cta_tile::K>
+>
+struct Gmem_tile_dout : public Base {
+
+    // Ctor.
+    template<typename Params, typename BInfo>
+    inline __device__ Gmem_tile_dout(const Params &params, const BInfo &binfo, int tidx)
+        : Base(params, 0, binfo, tidx) {
+
+        this->qkv_ptr_ = reinterpret_cast<char *>(params.o_ptr);
+        this->params_qkv_stride_in_bytes_ = params.o_stride_in_bytes;  // needed for move
+
+        // Compute the position of the thread in the row.
+        int col = tidx % Base::THREADS_PER_ROW;
+
+        // The row offset in the batched GEMM. For each seq element, we store O in that order.
+        int64_t row_offset = (int64_t)this->row_ * params.o_stride_in_bytes + binfo.bidx * Base::BYTES_PER_ROW;
+
+        // Assemble the final pointer.
+        this->qkv_ptr_ += row_offset + col * Base::BYTES_PER_LDG;
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Cta_tile, typename Base = fmha::Gmem_tile_o<Cta_tile> >
+struct Gmem_tile_dq : public Base {
+
+    // Ctor.
+    template<typename Params, typename BInfo>
+    inline __device__ Gmem_tile_dq(const Params &params, const BInfo &binfo, int tidx) 
+        : Base(params, binfo, tidx) {
+        this->o_ptr_ = reinterpret_cast<char *>(params.dqkv_ptr);
+        this->params_o_stride_in_bytes_ = params.qkv_stride_in_bytes;  // needed for move
+
+        // Compute the position of the thread in the row.
+        int col = tidx % Base::THREADS_PER_ROW;
+
+        // The row offset in the batched GEMM. For each seq element, we store O in that order.
+        int64_t row_offset = (int64_t)this->row_ * params.qkv_stride_in_bytes +
+                             (binfo.sum_s * 3 * binfo.h + binfo.bidh) * Base::BYTES_PER_ROW;
+
+        // Assemble the final pointer.
+        this->o_ptr_ += row_offset + col * Base::BYTES_PER_STG;
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace fmha
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/kernel_traits.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/kernel_traits.h
new file mode 100644
index 0000000000..965212070d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/kernel_traits.h
@@ -0,0 +1,95 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int S, int D, int STEP, int WARPS_M, int WARPS_N, uint32_t FLAGS = 0x8u>
+struct FMHA_kernel_traits {
+
+    // The CTA description for the 1st GEMM.
+    using Cta_tile_p = fmha::Cta_tile_extd<STEP, S, D, WARPS_M, WARPS_N, 1>;
+    // The CTA description for the 2nd GEMM.
+    using Cta_tile_o = fmha::Cta_tile_extd<STEP, D, S, WARPS_M, 1, WARPS_N>;
+
+    // Do we use one buffer for K and V.
+    enum { SHARE_SMEM_FOR_K_AND_V = (FLAGS & 0x8u) != 0u };
+
+    // The global memory tile to load Q.
+    using Gmem_tile_q = fmha::Gmem_tile_qkv<Cta_tile_p, fmha::BITS_PER_ELEMENT_A, STEP, D>;
+
+    // The shared memory tile to swizzle Q.
+    using Smem_tile_q = fmha::Smem_tile_a<Cta_tile_p, fmha::Row, Gmem_tile_q::BYTES_PER_LDG, 1>;
+
+    // The global memory tile to load K.
+    using Gmem_tile_k = fmha::Gmem_tile_qkv<Cta_tile_p, fmha::BITS_PER_ELEMENT_B, S, D>;
+    // The shared memory tile to swizzle K.
+    using Smem_tile_k = fmha::Smem_tile_b<Cta_tile_p, fmha::Col>;
+
+    // The global memory tile to load V.
+    using Gmem_tile_v = fmha::Gmem_tile_qkv<Cta_tile_o, fmha::BITS_PER_ELEMENT_B, S, D>;
+    // The shared memory tile to swizzle V.
+    using Smem_tile_v = fmha::Smem_tile_v<Cta_tile_o>;
+
+    // The global memory tile to store O.
+    using Gmem_tile_o = fmha::Gmem_tile_o<Cta_tile_o>;
+    // The shared memory tile for O.
+    using Smem_tile_o = fmha::Smem_tile_o<Cta_tile_o>;
+
+    // The global memory tile to load/store S.
+    using Gmem_tile_s = fmha::Gmem_tile_mma_s<Cta_tile_p>;
+
+    // The shared memory tile to transpose S.
+    using Smem_tile_st = fmha::Smem_tile_mma_transposed<Cta_tile_p>;
+
+    using Gmem_tile_do = fmha::Gmem_tile_dout<Cta_tile_p>;
+
+    // Make sure the number of threads match.
+    static_assert((int)Gmem_tile_o::THREADS_PER_ROW == (int)Smem_tile_o::THREADS_PER_ROW, "");
+
+    // The number of threads.
+    enum { THREADS = Cta_tile_p::THREADS_PER_CTA };
+    // Make sure the number of threads matches both CTAs.
+    static_assert((int)THREADS == (int)Cta_tile_o::THREADS_PER_CTA, "");
+
+    // The amount of shared memory needed to load Q and K.
+    enum { BYTES_PER_SMEM_QK = Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE };
+    // The extra amount of shared memory needed to load V.
+    enum { BYTES_PER_SMEM_V = SHARE_SMEM_FOR_K_AND_V ? 0u : Smem_tile_v::BYTES_PER_TILE };
+    // The amount of shared memory needed for Q, K and V..
+    enum { BYTES_PER_SMEM_QKV = BYTES_PER_SMEM_QK + BYTES_PER_SMEM_V };
+    // The amount of shared memory needed to load Q and store O.
+    enum { BYTES_PER_SMEM_QO = Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE };
+
+    // The amount of shared memory needed for Q, K, V and O.
+    enum { BYTES_PER_SMEM = fmha::Max<BYTES_PER_SMEM_QKV, BYTES_PER_SMEM_QO>::VALUE };
+    // Make sure we have enough shared memory.
+    static_assert(Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE <= BYTES_PER_SMEM, "");
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/mask.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/mask.h
new file mode 100644
index 0000000000..02db082215
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/mask.h
@@ -0,0 +1,76 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+namespace fmha {
+
+
+template<typename Cta_tile>
+struct Mask {
+    using Mma_tile = fmha::Hmma_tile<Cta_tile>;
+
+    template<typename Params, typename BInfo>
+    __device__ Mask(const Params &params, const BInfo &blockInfo, int tidx) {
+
+        actual_seqlen = blockInfo.actual_seqlen;
+
+        const int warp = tidx / Cta_tile::THREADS_PER_WARP;
+        const int lane = tidx % Cta_tile::THREADS_PER_WARP;
+
+        static_assert(Cta_tile::WARPS_K == 1, "");
+
+        // find the warp in the Cta tile
+        const int warp_n = (warp / Cta_tile::WARPS_M);
+        const int warp_m = (warp % Cta_tile::WARPS_M);
+        // decompose warp into 8x4 tile
+        const int quad = lane / 4;
+        const int tid = (lane % 4) * 2;
+        row = warp_m * 16 + quad;
+        col = warp_n * 16 + tid;
+    }
+
+    inline __device__ bool is_valid(const int mi, const int ni, const int ii, const int jj) const {
+
+        // ii and jj iterate over the 2x4 fragment
+        const bool col_valid = (ni * Mma_tile::N_PER_MMA_PER_CTA + col + (jj & 2) * 4 + (jj & 1)) < actual_seqlen;
+        //&& (row + mi * Mma_tile::M_PER_MMA_PER_CTA + ii * 8) < actual_seqlen;
+        return col_valid;
+        // return row_valid && col_valid;
+    }
+
+    inline __device__ void load(int it) {
+        row_offset = it * Cta_tile::M + row;
+    }
+    int row_offset;
+
+    int row;
+    int col;
+    int actual_seqlen;
+};
+
+}  // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/smem_tile.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/smem_tile.h
new file mode 100644
index 0000000000..7d7e0a5fcb
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/smem_tile.h
@@ -0,0 +1,1288 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include <fmha/utils.h>
+#include <fmha/gemm.h>
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< 
+    // The description of the tile computed by this CTA.
+    typename Cta_tile, 
+    // The number of rows in the 2D shared memory buffer.
+    int M_, 
+    // The number of cols.
+    int N_, 
+    // The size in bits of each element.
+    int BITS_PER_ELEMENT_, 
+    // The number of bytes per STS.
+    int BYTES_PER_STS_ = 16,
+    // The number of buffers. (Used in multistage and double buffer cases.)
+    int BUFFERS_PER_TILE_ = 1,
+    // Do we enable the fast path for LDS.128 and friends.
+    int ENABLE_LDS_FAST_PATH_ = 0, 
+    // The number of rows that are used for the XOR swizzling to allow fast STS/LDS. 
+    int ROWS_PER_XOR_PATTERN_ = 8,
+    // The number of cols that are used for the XOR swizzling to allow fast STS/LDS. 
+    int COLS_PER_XOR_PATTERN_ = 1,
+    // Use or not predicates
+    bool USE_PREDICATES_ = true
+>
+struct Smem_tile_without_skews {
+
+    // The size in bits of each element.
+    enum { BITS_PER_ELEMENT = BITS_PER_ELEMENT_ };
+    // The size in bytes of a single STS.
+    enum { BYTES_PER_STS = BYTES_PER_STS_ };
+    // The number of elements per STS.
+    enum { ELEMENTS_PER_STS = BYTES_PER_STS * 8 / BITS_PER_ELEMENT };
+    // To support arbitrary N, we pad some values to a power-of-2.
+    enum { N_WITH_PADDING = Next_power_of_two<N_>::VALUE }; 
+    // The number of bytes per row without packing of rows.
+    enum { BYTES_PER_ROW_BEFORE_PACKING = N_WITH_PADDING * BITS_PER_ELEMENT / 8 };
+    // The number of bytes per row -- we want at least 128B per row.
+    enum { BYTES_PER_ROW = Max<BYTES_PER_ROW_BEFORE_PACKING, 128>::VALUE };
+    // The number of rows in shared memory (two rows may be packed into a single one).
+    enum { ROWS = M_ * BYTES_PER_ROW_BEFORE_PACKING / BYTES_PER_ROW };
+
+    // The number of threads per row.
+    enum { THREADS_PER_ROW_UNBOUNDED = BYTES_PER_ROW / BYTES_PER_STS };
+    // The number of threads per row.
+    enum { THREADS_PER_ROW = Min<Cta_tile::THREADS_PER_CTA, THREADS_PER_ROW_UNBOUNDED>::VALUE };
+
+    // The number of STS per row.
+    enum { STS_PER_ROW = BYTES_PER_ROW / THREADS_PER_ROW / BYTES_PER_STS };
+    // It must be at least one.
+    static_assert(STS_PER_ROW >= 1, "");
+    // The number of rows written with a single STS.
+    enum { ROWS_PER_STS = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW };
+    // Make sure we write to at least one row per STS. Thanks Dr. Obvious ;)
+    static_assert(ROWS_PER_STS >= 1, "");
+    // The number of STS needed to store all rows.
+    enum { STS_PER_COL = Div_up<ROWS, ROWS_PER_STS>::VALUE };
+    // The number of STS in total.
+    enum { STS = STS_PER_COL * STS_PER_ROW };
+
+    // The size of one buffer in bytes in shared memory.
+    enum { BYTES_PER_BUFFER = STS * BYTES_PER_STS * Cta_tile::THREADS_PER_CTA };
+    // The number of buffers. 
+    enum { BUFFERS_PER_TILE = BUFFERS_PER_TILE_ };
+    // The size in bytes of total buffers.
+    enum { BYTES_PER_TILE = BYTES_PER_BUFFER * BUFFERS_PER_TILE };
+    // The boundary for smem_read_offset and smem_write_offset increment.
+    enum { BYTES_PER_TILE_INC_BOUNDARY = BYTES_PER_TILE - BYTES_PER_BUFFER };
+
+    // Do we enable the LDS.128 fast path?
+    enum { ENABLE_LDS_FAST_PATH = ENABLE_LDS_FAST_PATH_ };
+    static_assert(ENABLE_LDS_FAST_PATH == 0);
+    // The number of rows that are used for the XOR swizzling to allow fast STS/LDS. 
+    enum { ROWS_PER_XOR_PATTERN = ROWS_PER_XOR_PATTERN_ };
+    // The number of cols that are used for the XOR swizzling to allow fast STS/LDS. 
+    enum { COLS_PER_XOR_PATTERN = COLS_PER_XOR_PATTERN_ * 16 / BYTES_PER_STS };
+    // Use or not predicates
+    enum { USE_PREDICATES = USE_PREDICATES_ };
+
+    // The type of elements that are stored in shared memory by each thread.
+    using Store_type = typename Uint_from_size_in_bytes<BYTES_PER_STS>::Type;
+
+    // Ctor.
+    inline __device__ Smem_tile_without_skews(void *smem, int tidx) 
+        : smem_(__nvvm_get_smem_pointer(smem)) {
+
+        // The row written by a thread. See doc/mma_smem_layout.xlsx.
+        int smem_write_row = tidx / THREADS_PER_ROW;
+
+        // The XOR pattern.
+        int smem_write_xor = smem_write_row % ROWS_PER_XOR_PATTERN * COLS_PER_XOR_PATTERN;
+        // Compute the column and apply the XOR pattern.
+        int smem_write_col = (tidx % THREADS_PER_ROW) ^ smem_write_xor;
+
+        // The offset.
+        this->smem_write_offset_ = smem_write_row*BYTES_PER_ROW + smem_write_col*BYTES_PER_STS;
+
+        // TODO: Why not merge it with the read offset?
+        this->smem_read_buffer_ = __shfl_sync(0xffffffff, 0, 0);
+        this->smem_write_buffer_ = __shfl_sync(0xffffffff, 0, 0);
+    }
+
+    // Compute the store pointers.
+    template< int N >
+    inline __device__ void compute_store_pointers(uint32_t (&ptrs)[N]) {
+        #pragma unroll
+        for( int ii = 0; ii < N; ++ii ) {
+            // Decompose the STS into row/col.
+            int row = ii / STS_PER_ROW;
+            int col = ii % STS_PER_ROW;
+
+            // Assemble the offset.
+            int offset = smem_write_offset_ + row*ROWS_PER_STS*BYTES_PER_ROW;
+
+            // Take the column into account.
+            if( STS_PER_ROW > 1 ) {
+                offset += col*THREADS_PER_ROW*BYTES_PER_STS; 
+            }
+
+            // Apply the XOR pattern if needed.
+            if( ROWS_PER_STS < ROWS_PER_XOR_PATTERN ) {
+                const int m = row * ROWS_PER_STS % ROWS_PER_XOR_PATTERN;
+                offset ^= m * COLS_PER_XOR_PATTERN * BYTES_PER_STS;
+            }
+
+            // Assemble the final pointer :)
+            ptrs[ii] = smem_ + offset + smem_write_buffer_;
+        }
+    }
+
+    inline __device__ void debug_reset() {
+        for( int buffer = 0; buffer < BYTES_PER_TILE; buffer += BYTES_PER_BUFFER) {
+        for( int row = 0; row < ROWS; ++row ) {
+            for( int col = 0; col < BYTES_PER_ROW; col += 4 ) {
+                if( threadIdx.x == 0 ) {
+                    uint32_t val = 0x0;
+                    sts(val, smem_ + row*BYTES_PER_ROW + col + buffer);
+                }
+            }
+        }
+        }
+    }
+
+    // Print the content of the tile (only for debug ;)).
+    inline __device__ void debug_print() const {
+        for( int buffer = 0; buffer < BYTES_PER_TILE; buffer += BYTES_PER_BUFFER) {
+        for( int row = 0; row < ROWS; ++row ) {
+            for( int col = 0; col < BYTES_PER_ROW; col += 4 ) {
+                if( threadIdx.x == 0 ) {
+                    uint32_t val;
+                    lds(val, smem_ + row*BYTES_PER_ROW + col + buffer);
+                    printf("block=(x=%2d, y=%2d, z=%2d) (smem_=%2d, buffer=%2d, row=%2d, byte=%4d)=0x%08x\n",
+                        blockIdx.x,
+                        blockIdx.y,
+                        blockIdx.z,
+                        smem_,
+                        buffer,
+                        row,
+                        col,
+                        val);
+                }
+            }
+        }
+        }
+    }
+
+    // Move the read offset to next buffer.
+    inline __device__ void move_to_next_read_buffer() {
+        if( BUFFERS_PER_TILE > 1 && smem_read_buffer_ >= BYTES_PER_TILE_INC_BOUNDARY ) {
+            this->smem_read_buffer_ -= BYTES_PER_TILE_INC_BOUNDARY;
+        } else if( BUFFERS_PER_TILE > 1 ) {
+            this->smem_read_buffer_ += BYTES_PER_BUFFER;
+        }
+    }
+
+    // Move the read offset to next buffer. TODO: Remove this member function!!!
+    inline __device__ void move_next_read_buffer() {
+        this->move_to_next_read_buffer();
+    }
+
+    // Move the read offset to next N buffer (circular-buffer).
+    inline __device__ void move_to_next_read_buffer(int N) {
+        if( BUFFERS_PER_TILE > 1 ) {
+            this->smem_read_buffer_ += N * BYTES_PER_BUFFER;
+            this->smem_read_buffer_ -= smem_read_buffer_ >= BYTES_PER_TILE ? BYTES_PER_TILE : 0;
+        }
+    }
+
+    // Move the read offset to next N buffer (circular-buffer). TODO: Remove this member function!!!
+    inline __device__ void move_next_read_buffer(int N) {
+        this->move_to_next_read_buffer(N);
+    }
+
+    // Move the write offset to next buffer.
+    inline __device__ void move_to_next_write_buffer() {
+        if( BUFFERS_PER_TILE > 1 && smem_write_buffer_ >= BYTES_PER_TILE_INC_BOUNDARY ) {
+            this->smem_write_buffer_ -= BYTES_PER_TILE_INC_BOUNDARY;
+        } else if( BUFFERS_PER_TILE > 1 ) {
+            this->smem_write_buffer_ += BYTES_PER_BUFFER;
+        }
+    }
+
+    // Move the write offset to next buffer. TODO: Remove that member function!
+    inline __device__ void move_next_write_buffer() {
+        this->move_to_next_write_buffer();
+    }
+
+    // Move the read offset.
+    inline __device__ void move_read_offset(int delta) {
+        this->smem_read_offset_ += delta;
+    }
+
+    // Move the write offset.
+    inline __device__ void move_write_offset(int delta) {
+        this->smem_write_offset_ += delta;
+    }
+
+    // Store to the tile in shared memory.
+    template< int N >
+    inline __device__ void store(const Store_type (&data)[N], uint64_t = 0) {
+        uint32_t smem_ptrs[N];
+        this->compute_store_pointers(smem_ptrs);
+        sts(smem_ptrs, data);
+    }
+
+    // Store to the tile in shared memory.
+    template< int N, int M >
+    inline __device__ void store(const Store_type (&data)[N], uint32_t (&preds)[M], uint64_t = 0) {
+        uint32_t smem_ptrs[N];
+        this->compute_store_pointers(smem_ptrs);
+        sts(smem_ptrs, data, preds);
+    }
+
+    // Store to the tile in shared memory.
+    template< int N >
+    inline __device__ void store(const Store_type (&data)[N], uint32_t preds, uint64_t = 0) { 
+        this->store(data, preds);
+    }
+
+    // Store to the tile in shared memory.
+    template< int N >
+    inline __device__ void store(const void* (&gmem_ptrs)[N], uint32_t preds, uint64_t = 0) {
+        uint32_t tmp[1] = { preds };
+        this->store(gmem_ptrs, tmp);
+    }
+
+    // The shared memory pointer.
+    uint32_t smem_;
+    // The read offset. Reserve 4 offsets if needed.
+    int smem_read_offset_;
+    // The write offset.
+    int smem_write_offset_;
+    // The buffer base offset for read.
+    int smem_read_buffer_;
+    // The buffer base offset for write.
+    int smem_write_buffer_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< 
+    // The dimensions of the tile computed by the CTA.
+    typename Cta_tile, 
+    // The layout of the tile.
+    typename Layout, 
+    // The size of the STS.
+    int BYTES_PER_STS = 16,
+    // The number of buffers per tile.
+    int BUFFERS_PER_TILE = 1,
+    // Use or not predicates
+    bool USE_PREDICATES = true
+>
+struct Smem_tile_a {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int MMAS_K, int MMAS_K_WITH_PADDING >
+struct Compute_reset_mask {
+    // The potential mask.
+    enum { HALF = MMAS_K_WITH_PADDING / 2 };
+    // The remainder.
+    enum { MOD = MMAS_K % HALF };
+    // The final value.
+    enum { VALUE = (MMAS_K == MOD ? 0 : HALF) | Compute_reset_mask<MOD, HALF>::VALUE };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int MMAS_K_WITH_PADDING >
+struct Compute_reset_mask<0, MMAS_K_WITH_PADDING> {
+    enum { VALUE = 0 };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int MMAS_K >
+struct Compute_reset_mask<MMAS_K, MMAS_K> {
+    enum { VALUE = MMAS_K - 1 };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+struct Rows_per_xor_pattern_a {
+    // The size in bits.
+    enum { N_IN_BITS = N * fmha::BITS_PER_ELEMENT_A };
+    // The number of rows.
+    enum { VALUE = N_IN_BITS <= 256 ? 2 : (N_IN_BITS <= 512 ? 4 : 8) };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+struct Rows_per_xor_pattern_row_a : public Rows_per_xor_pattern_a<N> {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The dimensions of the tile computed by the CTA.
+    typename Cta_tile,
+    // The size of the STS.
+    int BYTES_PER_STS,
+    // The number of buffers per tile.
+    int BUFFERS_PER_TILE,
+    // How many rows to use for the XOR pattern to avoid bank conflicts?
+    int ROWS_PER_XOR_PATTERN_ = Rows_per_xor_pattern_row_a<Cta_tile::K>::VALUE
+>
+struct Smem_tile_row_a : public Smem_tile_without_skews<Cta_tile,
+                                                               Cta_tile::M,
+                                                               Cta_tile::K,
+                                                               fmha::BITS_PER_ELEMENT_A,
+                                                               BYTES_PER_STS,
+                                                               BUFFERS_PER_TILE,
+                                                               0,
+                                                               ROWS_PER_XOR_PATTERN_,
+                                                               1> {
+    // The MMA tile.
+    using Mma_tile = fmha::Hmma_tile<Cta_tile>;
+    // The base class.
+    using Base = Smem_tile_without_skews<Cta_tile,
+                                         Cta_tile::M,
+                                         Cta_tile::K,
+                                         fmha::BITS_PER_ELEMENT_A,
+                                         BYTES_PER_STS,
+                                         BUFFERS_PER_TILE,
+                                         0,
+                                         ROWS_PER_XOR_PATTERN_,
+                                         1>;
+    // The fragment.
+    using Fragment = Fragment_a<Row>;
+
+    // When we use padding to reach a power of two, special care has to be taken.
+    using Cta_tile_with_padding = Cta_tile_with_k_with_padding<Cta_tile>;
+    // The number of MMAs.
+    using Mma_tile_with_padding = fmha::Hmma_tile<Cta_tile_with_padding>;
+
+    // The size of a single LDS in bytes.
+    enum { BYTES_PER_LDS = 16 };
+
+    // Ctor.
+    inline __device__ Smem_tile_row_a(void *smem, int tidx) : Base(smem, tidx) {
+
+        // For documentation on the layout, see doc/mma_smem_layout.xlsx.
+
+        // The number of warps.
+        const int WARPS_M = Cta_tile::WARPS_M;
+        const int WARPS_N = Cta_tile::WARPS_N;
+        const int WARPS_K = Cta_tile::WARPS_K;
+
+        static_assert(WARPS_M == 1);
+        static_assert(WARPS_N == 4 || WARPS_N == 8);
+        static_assert(WARPS_K == 1);
+        static_assert(Base::ROWS_PER_XOR_PATTERN == 8);
+
+        // The row and column read by the thread.
+        int smem_read_row  = (tidx & 0x0f);
+        int smem_read_col  = (tidx & 0x07);
+        smem_read_col ^= (tidx & 0x10) / 16;
+
+        // The shared memory offset.
+        this->smem_read_offset_ = smem_read_row*Base::BYTES_PER_ROW + smem_read_col*BYTES_PER_LDS;
+    }
+
+    // Rewind smem_read_offset for last LDS phase in main loop.
+    inline __device__ void reverse_smem_read_offset(int ki = 0) {
+        // Undo the pointer increment for the next ni.
+        // Should match the load function below for ki = 0.
+        if( Mma_tile_with_padding::MMAS_K >=  2 ) {
+            this->smem_read_offset_ ^= BYTES_PER_LDS * 2;
+        }
+    }
+
+    // Load from shared memory.
+    inline __device__ void load(Fragment (&a)[Mma_tile::MMAS_M], int ki) {
+        #pragma unroll
+        for( int mi = 0; mi < Mma_tile::MMAS_M; ++mi ) {
+            // Jump by as many matrix rows as needed (a row in smem may pack multiple matrix rows).
+            int offset = mi * Mma_tile::M_PER_MMA_PER_CTA * Base::BYTES_PER_ROW_BEFORE_PACKING;
+
+            // Load using LDSM.M88.4.
+            uint4 tmp;
+            ldsm(tmp, this->smem_ + this->smem_read_offset_ + this->smem_read_buffer_ + offset);
+
+            // Store the value into the fragment.
+            a[mi].reg(0) = tmp.x;
+            a[mi].reg(1) = tmp.y;
+            a[mi].reg(2) = tmp.z;
+            a[mi].reg(3) = tmp.w;
+        }
+
+        // Move the offset to the next possition. See doc/mma_smem_layout.xlsx.
+        static_assert(Mma_tile_with_padding::MMAS_K < 64, "Not implemented");
+        if(        Mma_tile_with_padding::MMAS_K >= 32 && ki % 16 == 15 ) {
+            this->smem_read_offset_ ^= 31 * BYTES_PER_LDS * 2;
+        } else if( Mma_tile_with_padding::MMAS_K >= 16 && ki %  8 ==  7 ) {
+            this->smem_read_offset_ ^= 15 * BYTES_PER_LDS * 2;
+        } else if( Mma_tile_with_padding::MMAS_K >=  8 && ki %  4 ==  3 ) {
+            this->smem_read_offset_ ^=  7 * BYTES_PER_LDS * 2;
+        } else if( Mma_tile_with_padding::MMAS_K >=  4 && ki %  2 ==  1 ) {
+            this->smem_read_offset_ ^=  3 * BYTES_PER_LDS * 2;
+        } else if( Mma_tile_with_padding::MMAS_K >=  2 ) {
+            this->smem_read_offset_ ^=  1 * BYTES_PER_LDS * 2;
+        }
+    }
+
+    // Reset the read offset.
+    inline __device__ void reset_read_offset() {
+        // The number of MMAs in the K dimension.
+        enum { MMAS_K = Mma_tile::MMAS_K };
+        // The number of MMAs in the K dimension when we include padding.
+        enum { MMAS_K_WITH_PADDING = Mma_tile_with_padding::MMAS_K };
+        // Assemble the mask.
+        enum { MASK = Compute_reset_mask<MMAS_K, MMAS_K_WITH_PADDING>::VALUE };
+
+        // Reset the read offset.
+        this->smem_read_offset_ ^= MASK * BYTES_PER_LDS * 2;
+    }
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The dimensions of the tile computed by the CTA.
+    typename Cta_tile,
+    // The size of the STS.
+    int BYTES_PER_STS,
+    // The number of buffers per tile.
+    int BUFFERS_PER_TILE
+>
+struct Smem_tile_a<Cta_tile, Row, BYTES_PER_STS, BUFFERS_PER_TILE>
+    : public Smem_tile_row_a<Cta_tile,
+                                    BYTES_PER_STS,
+                                    BUFFERS_PER_TILE> {
+    // The base class.
+    using Base = Smem_tile_row_a<Cta_tile, BYTES_PER_STS, BUFFERS_PER_TILE>;
+
+    // Ctor.
+    inline __device__ Smem_tile_a(void *smem, int tidx) : Base(smem, tidx) {
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< 
+    // The dimensions of the tile computed by the CTA.
+    typename Cta_tile, 
+    // The layout of the tile.
+    typename Layout, 
+    // The size of the STS.
+    int BYTES_PER_STS = 16,
+    // The number of buffers per tile.
+    int BUFFERS_PER_TILE = 1,
+    // Use or not predicates
+    bool USE_PREDICATES = true
+>
+struct Smem_tile_b {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+struct Rows_per_xor_pattern_b {
+    // The size in bits.
+    enum { N_IN_BITS = N * fmha::BITS_PER_ELEMENT_B };
+    // The number of rows.
+    enum { VALUE = N_IN_BITS <= 256 ? 2 : (N_IN_BITS <= 512 ? 4 : 8) };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+struct Rows_per_xor_pattern_col_b : public Rows_per_xor_pattern_b<N> {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The dimensions of the tile computed by the CTA.
+    typename Cta_tile,
+    // The size of the STS.
+    int BYTES_PER_STS,
+    // The number of buffers per tile.
+    int BUFFERS_PER_TILE,
+    // How many rows to use for the XOR pattern to avoid bank conflicts?
+    int ROWS_PER_XOR_PATTERN_ = Rows_per_xor_pattern_col_b<Cta_tile::K>::VALUE
+>
+struct Smem_tile_col_b : public Smem_tile_without_skews<Cta_tile,
+                                                           Cta_tile::N,
+                                                           Cta_tile::K,
+                                                           fmha::BITS_PER_ELEMENT_B,
+                                                           BYTES_PER_STS,
+                                                           BUFFERS_PER_TILE,
+                                                           0,
+                                                           ROWS_PER_XOR_PATTERN_,
+                                                           1> {
+    // The MMA tile.
+    using Mma_tile = fmha::Hmma_tile<Cta_tile>;
+    // The base class.
+    using Base = Smem_tile_without_skews<Cta_tile,
+                                         Cta_tile::N,
+                                         Cta_tile::K,
+                                         fmha::BITS_PER_ELEMENT_B,
+                                         BYTES_PER_STS,
+                                         BUFFERS_PER_TILE,
+                                         0,
+                                         ROWS_PER_XOR_PATTERN_,
+                                         1>;
+    // The fragment.
+    using Fragment = Fragment_b< Col>;
+
+    // When we use padding to reach a power of two, special care has to be taken.
+    using Cta_tile_with_padding = Cta_tile_with_k_with_padding< Cta_tile>;
+    // The number of MMAs.
+    using Mma_tile_with_padding = fmha::Hmma_tile<Cta_tile_with_padding>;
+
+    // The size of a single LDS in bytes.
+    enum { BYTES_PER_LDS = 16 };
+
+    // The number of STS per thread
+    enum { STS_PER_THREAD_ = Base::ROWS * Base::THREADS_PER_ROW / Cta_tile::THREADS_PER_CTA };
+    // The number of STS per thread must be at least 1.
+    enum { STS_PER_THREAD = Max<1, STS_PER_THREAD_>::VALUE };
+
+    // Ctor.
+    inline __device__ Smem_tile_col_b(void *smem, int tidx) : Base(smem, tidx) {
+
+        // For documentation on the layout, see doc/mma_smem_layout.xlsx.
+
+        // The number of warps.
+        const int WARPS_M = Cta_tile::WARPS_M;
+        const int WARPS_N = Cta_tile::WARPS_N;
+        const int WARPS_K = Cta_tile::WARPS_K;
+        static_assert(Base::ROWS_PER_XOR_PATTERN == 8);
+        static_assert(WARPS_M == 1);
+        static_assert(WARPS_N == 4 || WARPS_N == 8);
+        static_assert(WARPS_K == 1);
+
+        // The masks to select the warps.
+        const int WARP_MASK_N = Warp_masks<WARPS_M, WARPS_N, WARPS_K>::N;
+
+        // The divisor for the warps.
+        const int WARP_DIV_N = WARPS_M *       1 * Cta_tile::THREADS_PER_WARP;
+
+        // The row and column read by the thread.
+        int smem_read_row  = (tidx & WARP_MASK_N) / WARP_DIV_N * Mma_tile::N_PER_MMA +
+                             (tidx & 0x07) +
+                             (tidx & 0x10) / 2;
+        int smem_read_col  = (tidx & 0x07);
+        smem_read_col ^= (tidx & 0x08) / 8;
+        // The shared memory offset.
+        this->smem_read_offset_ = smem_read_row*Base::BYTES_PER_ROW + smem_read_col*BYTES_PER_LDS;
+    }
+
+    // Rewind smem_read_offset for last LDS phase in main loop.
+    inline __device__ void reverse_smem_read_offset(int ki = 0) {
+        // Undo the pointer increment for the next ni.
+        // Should match the load function below for ki = 0.
+        if( Mma_tile_with_padding::MMAS_K >=  2 ) {
+            this->smem_read_offset_ ^= BYTES_PER_LDS * 2;
+        }
+    }
+
+    // Load from shared memory.
+    inline __device__ void load(Fragment (&b)[Mma_tile::MMAS_N], int ki) {
+        #pragma unroll
+        for( int ni = 0; ni < Mma_tile::MMAS_N; ++ni ) {
+            // Jump by as many matrix rows as needed (a row in smem may pack multiple matrix rows).
+            int offset = ni * Mma_tile::N_PER_MMA_PER_CTA * Base::BYTES_PER_ROW_BEFORE_PACKING;
+
+            // Load using LDSM.M88.4.
+            uint4 tmp;
+            ldsm(tmp, this->smem_ + this->smem_read_offset_ + this->smem_read_buffer_ + offset);
+
+            // Store the value into the fragment.
+            b[ni].reg(0) = tmp.x;
+            b[ni].reg(1) = tmp.y;
+            b[ni].reg(2) = tmp.z;
+            b[ni].reg(3) = tmp.w;
+        }
+
+        // Move the offset to the next possition. See doc/mma_smem_layout.xlsx.
+        static_assert(Mma_tile_with_padding::MMAS_K < 64, "Not implemented");
+        if(        Mma_tile_with_padding::MMAS_K >= 32 && ki % 16 == 15 ) {
+            this->smem_read_offset_ ^= 31 * BYTES_PER_LDS * 2;
+        } else if( Mma_tile_with_padding::MMAS_K >= 16 && ki %  8 ==  7 ) {
+            this->smem_read_offset_ ^= 15 * BYTES_PER_LDS * 2;
+        } else if( Mma_tile_with_padding::MMAS_K >=  8 && ki %  4 ==  3 ) {
+            this->smem_read_offset_ ^=  7 * BYTES_PER_LDS * 2;
+        } else if( Mma_tile_with_padding::MMAS_K >=  4 && ki %  2 ==  1 ) {
+            this->smem_read_offset_ ^=  3 * BYTES_PER_LDS * 2;
+        } else if( Mma_tile_with_padding::MMAS_K >=  2 ) {
+            this->smem_read_offset_ ^=  1 * BYTES_PER_LDS * 2;
+        }
+    }
+
+    // Reset the read offset.
+    inline __device__ void reset_read_offset() {
+        // The number of MMAs in the K dimension.
+        enum { MMAS_K = Mma_tile::MMAS_K };
+        // The number of MMAs in the K dimension when we include padding.
+        enum { MMAS_K_WITH_PADDING = Mma_tile_with_padding::MMAS_K };
+        // Assemble the mask.
+        enum { MASK = Compute_reset_mask<MMAS_K, MMAS_K_WITH_PADDING>::VALUE };
+
+        // Reset the read offset.
+        this->smem_read_offset_ ^= MASK * BYTES_PER_LDS * 2;
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The dimensions of the tile computed by the CTA.
+    typename Cta_tile,
+    // The size of the STS.
+    int BYTES_PER_STS,
+    // The number of buffers per tile.
+    int BUFFERS_PER_TILE
+>
+struct Smem_tile_b< Cta_tile, Col, BYTES_PER_STS, BUFFERS_PER_TILE >
+    : public Smem_tile_col_b<Cta_tile,
+                             BYTES_PER_STS,
+                             BUFFERS_PER_TILE> {
+
+    // The base class.
+    using Base = Smem_tile_col_b< Cta_tile, BYTES_PER_STS, BUFFERS_PER_TILE>;
+
+    // Ctor.
+    inline __device__ Smem_tile_b(void *smem, int tidx) : Base(smem, tidx) {
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<  int N >
+struct Rows_per_xor_pattern_row_b : public Rows_per_xor_pattern_b< N> {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+template<
+    // The dimensions of the tile computed by the CTA.
+    typename Cta_tile,
+    // The size of the STS.
+    int BYTES_PER_STS,
+    // The number of buffers per tile.
+    int BUFFERS_PER_TILE,
+    // How many rows to use for the XOR pattern to avoid bank conflicts?
+    int ROWS_PER_XOR_PATTERN_ = Rows_per_xor_pattern_row_b<Cta_tile::N>::VALUE,
+    // How many cols to use for the XOR pattern to avoid bank conflicts?
+    int COLS_PER_XOR_PATTERN_ = 1
+>
+struct Smem_tile_row_b : public Smem_tile_without_skews<Cta_tile,
+                                                               Cta_tile::K,
+                                                               Cta_tile::N,
+                                                               fmha::BITS_PER_ELEMENT_B,
+                                                               BYTES_PER_STS,
+                                                               BUFFERS_PER_TILE,
+                                                               0,
+                                                               ROWS_PER_XOR_PATTERN_,
+                                                               COLS_PER_XOR_PATTERN_> {
+
+    // The MMA tile.
+    using Mma_tile = fmha::Hmma_tile<Cta_tile>;
+    // The base class.
+    using Base = Smem_tile_without_skews<Cta_tile,
+                                         Cta_tile::K,
+                                         Cta_tile::N,
+                                         fmha::BITS_PER_ELEMENT_B,
+                                         BYTES_PER_STS,
+                                         BUFFERS_PER_TILE,
+                                         0,
+                                         ROWS_PER_XOR_PATTERN_,
+                                         COLS_PER_XOR_PATTERN_>;
+    // The fragment.
+    using Fragment = Fragment_b<Row>;
+
+    // Can we use LDSM? No if the data type is 32-bit large.
+    enum { USE_LDSMT = fmha::BITS_PER_ELEMENT_B == 16 };
+    // The size of a single LDS in bytes.
+    enum { BYTES_PER_LDS = USE_LDSMT ? 16 : 4 };
+    // The number of elements per LDS.
+    enum { ELEMENTS_PER_LDS = BYTES_PER_LDS * 8 / fmha::BITS_PER_ELEMENT_B };
+
+    // The number of STS per thread
+    enum { STS_PER_THREAD_ = Base::ROWS * Base::THREADS_PER_ROW / Cta_tile::THREADS_PER_CTA };
+    // The number of STS per thread must be at least 1.
+    enum { STS_PER_THREAD = Max<1, STS_PER_THREAD_>::VALUE };
+
+    // Ctor.
+    inline __device__ Smem_tile_row_b(void *smem, int tidx) : Base(smem, tidx) {
+
+        // The number of warps.
+        const int WARPS_M = Cta_tile::WARPS_M;
+        const int WARPS_N = Cta_tile::WARPS_N;
+        const int WARPS_K = Cta_tile::WARPS_K;
+        static_assert(WARPS_K == 1);
+        static_assert(WARPS_M == 4 || WARPS_M == 8);
+        static_assert(WARPS_N == 1);
+
+        // The masks to select the warps.
+        const int WARP_MASK_N = Warp_masks<WARPS_M, WARPS_N, WARPS_K>::N;
+        const int WARP_MASK_K = Warp_masks<WARPS_M, WARPS_N, WARPS_K>::K;
+
+        // The divisor for the warps.
+        const int WARP_DIV_N = WARPS_M *       1 * Cta_tile::THREADS_PER_WARP;
+        const int WARP_DIV_K = WARPS_M * WARPS_N * Cta_tile::THREADS_PER_WARP;
+
+        // The row/col read by the thread.
+        int smem_read_row, smem_read_col;
+
+        static_assert(USE_LDSMT);
+        static_assert(Base::ROWS_PER_XOR_PATTERN == 8);
+
+        smem_read_row = (tidx & WARP_MASK_K) / WARP_DIV_K * Mma_tile::MMAS_K * 16 +
+                        (tidx & 0x07) + (tidx & 0x08);
+        smem_read_col = (tidx & 0x07);
+        smem_read_col ^= (tidx & WARP_MASK_N) / WARP_DIV_N * 2 + (tidx & 0x10) / 16;
+
+        // The shared memory offset.
+        this->smem_read_offset_ = smem_read_row*Base::BYTES_PER_ROW + smem_read_col*BYTES_PER_LDS;
+
+        // Fill zeroes for group conv
+    }
+
+    // Rewind smem_read_offset for last LDS phase in main loop.
+    inline __device__ void reverse_smem_read_offset(int ki = 0) {
+        // The size of each element in bits.
+        const int BITS_PER_ELT = fmha::BITS_PER_ELEMENT_B;
+        // The size in bytes of the data needed to compute an MMA per CTA.
+        const int BYTES_PER_MMA_PER_CTA = Mma_tile::N_PER_MMA_PER_CTA * BITS_PER_ELT / 8;
+
+        #pragma unroll
+        for( int ni = 0; ni < Mma_tile::MMAS_N; ++ni ) {
+            // Undo the pointer increment for the next ni.
+            // Should match the load function below for ki = 0.
+            if( BYTES_PER_MMA_PER_CTA >= 128 ) {
+                // Nothing to do!
+            } else if( BYTES_PER_MMA_PER_CTA == 64 && Mma_tile::MMAS_N > 1 ) {
+                this->smem_read_offset_ ^= BYTES_PER_MMA_PER_CTA;
+            } else if( BYTES_PER_MMA_PER_CTA == 64 ) {
+                // Nothing to do!
+            } else if( BYTES_PER_MMA_PER_CTA == 32 && Mma_tile::MMAS_N == 4 ) {
+                this->smem_read_offset_ ^= BYTES_PER_LDS * (ni % 2 == 0 ? 2 : 6);
+            } else if( BYTES_PER_MMA_PER_CTA == 32 && Mma_tile::MMAS_N == 2 ) {
+                this->smem_read_offset_ ^= BYTES_PER_LDS * 2;
+            }
+        }
+
+        // Reset smem_read_offset for odd MMAS_N > 1 (npo2 kernels)
+        if( BYTES_PER_MMA_PER_CTA == 64 && Mma_tile::MMAS_N > 1 &&
+                Mma_tile::MMAS_N % 2 == 1 ) {
+            this->smem_read_offset_ ^= BYTES_PER_MMA_PER_CTA;
+        }
+    }
+
+    // Load from shared memory.
+    inline __device__ void load(Fragment (&b)[Mma_tile::MMAS_N], int ki) {
+        // The size of each element in bits.
+        const int BITS_PER_ELT = fmha::BITS_PER_ELEMENT_B;
+        // The size in bytes of the data needed to compute an MMA per CTA.
+        const int BYTES_PER_MMA_PER_CTA = Mma_tile::N_PER_MMA_PER_CTA * BITS_PER_ELT / 8;
+
+        #pragma unroll
+        for( int ni = 0; ni < Mma_tile::MMAS_N; ++ni ) {
+            // Prepare the offset.
+            int offset = ki * Base::ROWS_PER_XOR_PATTERN * 2 * Base::BYTES_PER_ROW;
+                if ( BYTES_PER_MMA_PER_CTA == 32 ) {
+                    offset += this->smem_read_offset_;
+                } else if ( BYTES_PER_MMA_PER_CTA == 64 ) {
+                    offset += this->smem_read_offset_ + (ni/2) * BYTES_PER_MMA_PER_CTA * 2;
+                } else {
+                    offset += this->smem_read_offset_ + (ni  ) * BYTES_PER_MMA_PER_CTA;
+                }
+
+            // Load the data using LDSM.MT88.2.
+            uint32_t ptr = this->smem_ + this->smem_read_buffer_ + offset;
+            uint4 tmp;
+            if( USE_LDSMT ) {
+                ldsmt(tmp, ptr);
+            } else {
+                lds(tmp.x, (ptr     ) + 0*Base::BYTES_PER_ROW);
+                lds(tmp.y, (ptr     ) + 4*Base::BYTES_PER_ROW);
+                lds(tmp.z, (ptr ^ 32) + 0*Base::BYTES_PER_ROW);
+                lds(tmp.w, (ptr ^ 32) + 4*Base::BYTES_PER_ROW);
+            }
+
+            // Store those values in the fragment.
+            b[ni].reg(0) = tmp.x;
+            b[ni].reg(1) = tmp.y;
+            b[ni].reg(2) = tmp.z;
+            b[ni].reg(3) = tmp.w;
+
+            // Move the pointer for the next ni. I expect the compiler to not recompute those.
+            if( BYTES_PER_MMA_PER_CTA >= 128 ) {
+                // Nothing to do!
+            } else if( BYTES_PER_MMA_PER_CTA == 64 && Mma_tile::MMAS_N > 1 ) {
+                this->smem_read_offset_ ^= BYTES_PER_MMA_PER_CTA;
+            } else if( BYTES_PER_MMA_PER_CTA == 64 ) {
+                // Nothing to do!
+            } else if( BYTES_PER_MMA_PER_CTA == 32 && Mma_tile::MMAS_N == 4 ) {
+                this->smem_read_offset_ ^= BYTES_PER_LDS * (ni % 2 == 0 ? 2 : 6);
+            } else if( BYTES_PER_MMA_PER_CTA == 32 && Mma_tile::MMAS_N == 2 ) {
+                this->smem_read_offset_ ^= BYTES_PER_LDS * 2;
+            }
+        }
+
+        // Reset smem_read_offset for odd MMAS_N > 1 (npo2 kernels)
+        if( BYTES_PER_MMA_PER_CTA == 64 && Mma_tile::MMAS_N > 1 &&
+                Mma_tile::MMAS_N % 2 == 1 ) {
+            this->smem_read_offset_ ^= BYTES_PER_MMA_PER_CTA;
+        }
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    // The dimensions of the tile computed by the CTA.
+    typename Cta_tile,
+    // The size of the STS.
+    int BYTES_PER_STS,
+    // The number of buffers per tile.
+    int BUFFERS_PER_TILE
+>
+struct Smem_tile_b<Cta_tile, Row, BYTES_PER_STS, BUFFERS_PER_TILE>
+    : public Smem_tile_row_b<Cta_tile,
+                             BYTES_PER_STS,
+                             BUFFERS_PER_TILE> {
+
+    // The base class.
+    using Base = Smem_tile_row_b<Cta_tile, BYTES_PER_STS, BUFFERS_PER_TILE>;
+
+    // Ctor.
+    inline __device__ Smem_tile_b(void *smem, int tidx) : Base(smem, tidx) {
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Cta_tile>
+struct Smem_tile_v : public fmha::Smem_tile_without_skews<Cta_tile, Cta_tile::K, Cta_tile::N, 16, 16, 1, 0, 8, 1> {
+
+    // The base class.
+    using Base = Smem_tile_without_skews<Cta_tile, Cta_tile::K, Cta_tile::N, 16, 16, 1, 0, 8, 1>;
+    // The MMA tile.
+    using Mma_tile = fmha::Hmma_tile<Cta_tile>;
+    // The fragment.
+    using Fragment = Fragment_b< fmha::Col>;
+
+    // The size of a single LDS in bytes.
+    enum { BYTES_PER_LDS = 16 };
+
+    // Ctor.
+    inline __device__ Smem_tile_v(void *smem, int tidx) : Base(smem, tidx) {
+
+        // The row/col read by the thread.
+        int read_row, read_col;
+
+        static_assert(Cta_tile::WARPS_M == 1 && Cta_tile::WARPS_N == 1 && (Cta_tile::WARPS_K == 4 || Cta_tile::WARPS_K == 8));
+
+        read_row = (tidx & 0xe0) / 2 + (tidx & 0x0f);
+        read_col = (tidx & 0x07);
+        read_col ^= (tidx & 0x10) / 16;
+
+        // The shared memory offset.
+        this->smem_read_offset_ = read_row * Base::BYTES_PER_ROW + read_col * BYTES_PER_LDS;
+    }
+
+    // Load from shared memory.
+    inline __device__ void load(Fragment (&b)[Mma_tile::MMAS_N], int ki) {
+#pragma unroll
+        for( int ni = 0; ni < Mma_tile::MMAS_N; ++ni ) {
+            // Jump by 16 * #warps row.
+            int row = ki * 16 * Cta_tile::WARPS_K;
+
+            // Load the data using LDSM.MT88.2.
+            uint4 tmp;
+            fmha::ldsmt(tmp, this->smem_ + this->smem_read_offset_ + row * Base::BYTES_PER_ROW);
+            b[ni].reg(0) = tmp.x;
+            b[ni].reg(1) = tmp.y;
+            b[ni].reg(2) = tmp.z;
+            b[ni].reg(3) = tmp.w;
+
+            // Move the pointer for the next ni. I expect the compiler to not recompute those.
+            if( Mma_tile::MMAS_N == 4 ) {
+                this->smem_read_offset_ ^= BYTES_PER_LDS * (ni % 2 == 0 ? 2 : 6);
+            } else {
+                assert(false);  // Not implemented!
+            }
+        }
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Cta_tile>
+struct Smem_tile_o {
+
+    // The MMA tile.
+    using Mma_tile = fmha::Hmma_tile<Cta_tile>;
+    // The accumulators.
+    using Accumulator = fmha::Fragment_accumulator;
+    // The accumulators.
+    using Data_type = typename Accumulator::Data_type;
+
+    // The size of each element.
+    enum { BYTES_PER_ELEMENT = sizeof(Data_type) };
+    // The size of each STS.
+    enum { BYTES_PER_STS = 8 };
+    // The size of each row in shared memory.
+    enum { BYTES_PER_ROW = Cta_tile::N * Cta_tile::WARPS_K * BYTES_PER_ELEMENT };
+
+    // The size of each LDS.
+    enum { BYTES_PER_LDS = 16 };
+    enum { THREADS_PER_ROW = 16 };
+
+    // The number of rows.
+    enum { ROWS = Cta_tile::M };
+    // The number of "rows" to process per loop iteration (in the "epilogue").
+    enum { ROWS_PER_LOOP = ROWS <= 64 ? ROWS : (int)Mma_tile::M_PER_MMA_PER_CTA };
+    // The number of outer loops.
+    enum { LOOPS = ROWS / ROWS_PER_LOOP };
+    // Make sure it matches our expectations.
+    static_assert(LOOPS == 1 || LOOPS == (int)Mma_tile::MMAS_M, "");
+
+    // The number of rows loaded per LDS.
+    enum { ROWS_PER_LDS = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW };
+    // Do we have to guard against partial writes/reads.
+    enum { HAS_INCOMPLETE_LDS = ROWS_PER_LOOP % ROWS_PER_LDS != 0 };
+    // The total number of LDS per loop.
+    enum { LDS_PER_LOOP = fmha::Div_up<ROWS_PER_LOOP, ROWS_PER_LDS>::VALUE };
+
+    // The amount of shared memory.
+    enum { BYTES_PER_TILE = ROWS_PER_LOOP * BYTES_PER_ROW };
+
+    // The write pointer.
+    uint32_t smem_write_, smem_read_;
+    // Is the thread active for the last LDS of the series?
+    int is_active_for_last_lds_;
+
+    static_assert(BYTES_PER_ROW == 64 * 4 * Cta_tile::WARPS_K);
+    static_assert(LOOPS == 1 || LOOPS == (int)Mma_tile::MMAS_M, "");
+
+    // Ctor.
+    inline __device__ Smem_tile_o(void *smem, int tidx) {
+
+        // Get a 32-bit value for the shared memory address.
+        uint32_t smem_ = __nvvm_get_smem_pointer(smem);
+
+        static_assert(Cta_tile::WARPS_M == 1 && Cta_tile::WARPS_N == 1 && (Cta_tile::WARPS_K == 4 || Cta_tile::WARPS_K == 8));
+
+        int write_row = (tidx & 0x1c) / 4;
+        int write_col = (tidx);
+
+        // Assemble the write pointer.
+        smem_write_ = smem_ + write_row * BYTES_PER_ROW + write_col * BYTES_PER_STS;
+
+        // The element read by each thread.
+        int read_row = tidx / THREADS_PER_ROW;
+        int read_col = tidx % THREADS_PER_ROW;
+
+        // Take the XOR pattern into account for the column.
+        read_col ^= 2 * (read_row & 0x7);
+
+        // Assemble the read pointer.
+        this->smem_read_ = smem_ + read_row * BYTES_PER_ROW + read_col * BYTES_PER_LDS;
+
+        // Is that thread active on the last LDS?
+        if( HAS_INCOMPLETE_LDS ) {
+            this->is_active_for_last_lds_ = read_row + (LDS_PER_LOOP - 1) * ROWS_PER_LDS < Cta_tile::M;
+        }
+    }
+
+    // Load the output fragments.
+    inline __device__ void load(uint4 (&out)[LDS_PER_LOOP]) const {
+        #pragma unroll
+        for( int ii = 0; ii < LDS_PER_LOOP; ++ii ) {
+
+            // Load the elements before the reduction (split-K).
+            uint4 tmp[Cta_tile::WARPS_K];
+            #pragma unroll
+            for( int jj = 0; jj < Cta_tile::WARPS_K; ++jj ) {
+                int imm = ii * ROWS_PER_LDS * BYTES_PER_ROW + jj * Cta_tile::N * BYTES_PER_ELEMENT;
+                if( !HAS_INCOMPLETE_LDS || (ii < LDS_PER_LOOP - 1 || this->is_active_for_last_lds_) ) {
+                    fmha::lds(tmp[jj], this->smem_read_ + imm);
+                }
+            }
+
+            // Perform the reduction.
+            out[ii] = tmp[0];
+            #pragma unroll
+            for( int jj = 1; jj < Cta_tile::WARPS_K; ++jj ) {
+                out[ii] = fmha::fadd4(out[ii], tmp[jj]);
+            }
+        }
+    }
+    // Store the accumulators.
+    template <int M, int N>
+    inline __device__ void store(const Accumulator (&acc)[M][N], int mi) {
+        enum { M_PER_MMA = Mma_tile::M_PER_MMA_PER_CTA };
+        #pragma unroll
+        for( int ni = 0; ni < Mma_tile::MMAS_N; ++ni ) {
+
+            // The number of MMAs that are stored per loop iteration.
+            enum { MMAS_M_PER_LOOP = Mma_tile::MMAS_M / LOOPS };
+
+            // Store 1st column of the different MMAs.
+            #pragma unroll
+            for( int mj = 0; mj < MMAS_M_PER_LOOP; ++mj ) {
+                // Precompute the immediates to jump between rows.
+                int row_0 = (mj * M_PER_MMA + 0) * BYTES_PER_ROW;
+                int row_1 = (mj * M_PER_MMA + 8) * BYTES_PER_ROW;
+                uint2 tmp0, tmp1;
+                tmp0.x = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(0);
+                tmp0.y = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(1);
+
+                tmp1.x = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(2);
+                tmp1.y = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(3);
+
+                // Store.
+                fmha::sts(this->smem_write_ + row_0, tmp0);
+                fmha::sts(this->smem_write_ + row_1, tmp1);
+            }
+
+            // Swizzle the write pointer using a XOR of 16B.
+            this->smem_write_ ^= 32;
+
+            // Store 2nd column of the different MMAs.
+            #pragma unroll
+            for( int mj = 0; mj < MMAS_M_PER_LOOP; ++mj ) {
+                // Precompute the immediates to jump between rows.
+                int row_0 = (mj * M_PER_MMA + 0) * BYTES_PER_ROW;
+                int row_1 = (mj * M_PER_MMA + 8) * BYTES_PER_ROW;
+
+                uint2 tmp0, tmp1;
+                tmp0.x = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(4);
+                tmp0.y = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(5);
+
+                tmp1.x = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(6);
+                tmp1.y = acc[mi * MMAS_M_PER_LOOP + mj][ni].reg(7);
+                // Store.
+                fmha::sts(this->smem_write_ + row_0, tmp0);
+                fmha::sts(this->smem_write_ + row_1, tmp1);
+            }
+
+            // Cancel the previous XOR of 1 + swizzle the write pointer using a XOR of 32B or 64B.
+            this->smem_write_ ^= (ni & 1) ? 7 * 32 : 3 * 32;
+        }
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Cta_tile>
+struct Smem_tile_mma {
+
+    using Mma_tile = fmha::Hmma_tile<Cta_tile>;
+    using Fragment = fmha::Fragment_a<fmha::Col>;
+
+    enum { COLS = Cta_tile::N };
+    enum { BYTES_PER_ELT = 2 };
+    enum { BYTES_PER_STS = 4 };
+    enum { BYTES_PER_ROW = COLS * BYTES_PER_ELT };  // TODO
+    enum { BYTES_PER_TILE = Cta_tile::M * BYTES_PER_ROW };
+
+    enum { WARPS_M = Cta_tile::WARPS_M };
+    enum { WARPS_N = Cta_tile::WARPS_N };
+    enum { WARPS_K = Cta_tile::WARPS_K };
+
+    static_assert(WARPS_K == 1);
+    inline __device__ Smem_tile_mma(char *smem, int tidx) {
+        smem_ = __nvvm_get_smem_pointer(smem);
+
+        int write_col, write_row;
+        static_assert(WARPS_M == 1 && (WARPS_N == 4 || WARPS_N == 8) || (WARPS_M == 4 || WARPS_N == 8) || WARPS_N == 1);
+        if( WARPS_M == 1 && (WARPS_N == 4 || WARPS_N == 8) ) {
+            write_row = (tidx & 0x1c) / 4;
+            write_col = (tidx & 0xe0) / 4 + (tidx & 0x03);
+        } else {
+            write_row = (tidx & 0xe0) / 2 + (tidx & 0x1c) / 4;
+            write_col = (tidx & 0x03);
+        }
+        write_col ^= (write_row & 0x07) * 4;
+
+        write_offset_ = write_row * BYTES_PER_ROW + write_col * BYTES_PER_STS;
+    }
+
+    template<int M, int N>
+    inline __device__ void store(const uint4 (&regs)[M][N]) {
+        static_assert(COLS == Cta_tile::N);
+        for( int mi = 0; mi < M; mi++ ) {
+            for( int ni = 0; ni < N; ni++ ) {
+                size_t offset = write_offset_ + mi * WARPS_M * 16 * BYTES_PER_ROW + ni * WARPS_N * 16 * BYTES_PER_ELT;
+                fmha::sts(smem_ + offset + 0 * BYTES_PER_ROW, regs[mi][ni].x);
+                fmha::sts(smem_ + offset + 8 * BYTES_PER_ROW, regs[mi][ni].z);
+                offset ^= 4 * BYTES_PER_STS;
+                fmha::sts(smem_ + offset + 0 * BYTES_PER_ROW, regs[mi][ni].y);
+                fmha::sts(smem_ + offset + 8 * BYTES_PER_ROW, regs[mi][ni].w);
+            }
+        }
+    }
+
+    uint32_t smem_;
+    uint32_t write_offset_;
+    uint32_t warp_m;
+    uint32_t warp_n;
+    uint32_t lane;
+};
+
+template< typename Cta_tile, typename Base = Smem_tile_mma< Cta_tile>>
+struct Smem_tile_mma_transposed : public Base {
+    enum { BYTES_PER_LDS = 16 };
+    enum { BYTES_PER_ROW = Base::BYTES_PER_ROW };
+    enum { BYTES_PER_ELT = Base::BYTES_PER_ELT };
+    enum { WARPS_M = Base::WARPS_M };
+    enum { WARPS_N = Base::WARPS_N };
+    static_assert(WARPS_M == 1 && (WARPS_N == 4 || WARPS_N == 8));
+    using Fragment = typename Base::Fragment;
+    inline __device__ Smem_tile_mma_transposed(char *smem, int tidx) : Base(smem, tidx) {
+
+        static_assert(WARPS_M == 1 && (WARPS_N == 4 || WARPS_N == 8));
+        int read_row, read_col;
+        read_row = (tidx & 0x0f);
+        read_col = (tidx & 0xe0) / 16 + (tidx & 0x1c) / 16;
+
+        read_col ^= (read_row & 0x07);
+        read_offset_ = read_row * BYTES_PER_ROW + read_col * BYTES_PER_LDS;
+    }
+
+    template<int M, int N>
+    inline __device__ void load(Fragment (&frag)[M][N]) {
+        static_assert(Base::COLS == Cta_tile::N);
+        for( int mi = 0; mi < M; mi++ ) {
+            for( int ni = 0; ni < N; ni++ ) {
+                size_t offset = read_offset_ + mi * WARPS_M * 16 * BYTES_PER_ROW + ni * WARPS_N * 16 * BYTES_PER_ELT;
+                uint4 dst;
+                fmha::ldsmt(dst, this->smem_ + offset);
+                frag[mi][ni].reg(0) = dst.x;
+                frag[mi][ni].reg(1) = dst.z;  // Fragment A regs col major!
+                frag[mi][ni].reg(2) = dst.y;
+                frag[mi][ni].reg(3) = dst.w;
+            }
+        }
+    }
+
+    uint32_t read_offset_;
+};
+
+template< typename Cta_tile, typename Base = Smem_tile_mma< Cta_tile>>
+struct Smem_tile_mma_epilogue : public Base {
+    enum { BYTES_PER_LDS = 16 };
+    enum { BYTES_PER_ROW = Base::BYTES_PER_ROW };
+    enum { BYTES_PER_ELT = Base::BYTES_PER_ELT };
+    enum { THREADS_PER_ROW = BYTES_PER_ROW / BYTES_PER_LDS };
+    static_assert(THREADS_PER_ROW * BYTES_PER_LDS == BYTES_PER_ROW);
+    enum { ROWS_PER_LDS = Cta_tile::THREADS_PER_CTA / THREADS_PER_ROW };
+    enum { NUM_LDS = Cta_tile::M / ROWS_PER_LDS };
+    static_assert(NUM_LDS * ROWS_PER_LDS == Cta_tile::M);
+    enum { WARPS_M = Base::WARPS_M };
+    enum { WARPS_N = Base::WARPS_N };
+    static_assert((WARPS_M == 4 || WARPS_N == 8) || WARPS_N == 1);
+    
+    using Acc = fmha::Fragment_accumulator;
+
+    inline __device__ Smem_tile_mma_epilogue(char *smem, int tidx) : Base(smem, tidx) {
+        const int read_row = tidx / THREADS_PER_ROW;
+        int read_col = tidx % THREADS_PER_ROW;
+        read_col ^= (read_row & 0x07);
+        read_offset_ = read_row * BYTES_PER_ROW + read_col * BYTES_PER_LDS;
+    }
+
+    inline __device__ void load(uint4 (&data)[NUM_LDS]) {
+        for( int ii = 0; ii < NUM_LDS; ii++ ) {
+            size_t offset = read_offset_ + ii * ROWS_PER_LDS * BYTES_PER_ROW;
+            fmha::lds(data[ii], this->smem_ + offset);
+        }
+    }
+
+    template<int M, int N>
+    inline __device__ void store(const Acc (&acc)[M][N]){
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ni = 0; ni < N; ni++ ) {
+                // 1st row - 4 elements per row.
+                float tmp00 = acc[mi][ni].elt(0);
+                float tmp01 = acc[mi][ni].elt(1);
+                float tmp02 = acc[mi][ni].elt(4);
+                float tmp03 = acc[mi][ni].elt(5);
+                // 2nd row - 4 elements per row.
+                float tmp10 = acc[mi][ni].elt(2);
+                float tmp11 = acc[mi][ni].elt(3);
+                float tmp12 = acc[mi][ni].elt(6);
+                float tmp13 = acc[mi][ni].elt(7);
+
+                uint32_t x = fmha::float2_to_half2(tmp00, tmp01);
+                uint32_t y = fmha::float2_to_half2(tmp02, tmp03);
+                uint32_t z = fmha::float2_to_half2(tmp10, tmp11);
+                uint32_t w = fmha::float2_to_half2(tmp12, tmp13);
+     
+                size_t offset = (this->write_offset_ ^ (ni * 32)) + mi * WARPS_M * 16 * BYTES_PER_ROW;
+                fmha::sts(this->smem_ + offset + 0 * BYTES_PER_ROW, x);
+                fmha::sts(this->smem_ + offset + 8 * BYTES_PER_ROW, z);
+                offset ^= 4 * Base::BYTES_PER_STS;
+                fmha::sts(this->smem_ + offset + 0 * BYTES_PER_ROW, y);
+                fmha::sts(this->smem_ + offset + 8 * BYTES_PER_ROW, w);
+            }
+        }
+    }
+
+
+
+    template<int M, int N>
+    inline __device__ void store(const uint4 (&regs)[M][N]) {
+        for( int mi = 0; mi < M; mi++ ) {
+            for( int ni = 0; ni < N; ni++ ) {
+                size_t offset = (this->write_offset_ ^ (ni * 32)) + mi * WARPS_M * 16 * BYTES_PER_ROW;
+                fmha::sts(this->smem_ + offset + 0 * BYTES_PER_ROW, regs[mi][ni].x);
+                fmha::sts(this->smem_ + offset + 8 * BYTES_PER_ROW, regs[mi][ni].z);
+                offset ^= 4 * Base::BYTES_PER_STS;
+                fmha::sts(this->smem_ + offset + 0 * BYTES_PER_ROW, regs[mi][ni].y);
+                fmha::sts(this->smem_ + offset + 8 * BYTES_PER_ROW, regs[mi][ni].w);
+            }
+        }
+    }
+
+    uint32_t read_offset_;
+};
+
+}  // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/softmax.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/softmax.h
new file mode 100644
index 0000000000..22bf1abdfc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/softmax.h
@@ -0,0 +1,478 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Sum_ {
+    enum { IS_SUM = 1 };
+    static inline __device__ float apply(float x, float y) {
+        return x + y;
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Max_ {
+    enum { IS_SUM = 0 };
+    static inline __device__ float apply(float x, float y) {
+        return x > y ? x : y;
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float apply_exp_(float x, float max) {
+    return __expf(x - max);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Cta_tile, typename Kernel_traits>
+struct Softmax_base {
+
+    // The Mma tile.
+    using Mma_tile = fmha::Hmma_tile<Cta_tile>;
+
+    // The number of MMAs in M/N dimensions.
+    enum { MMAS_M = Mma_tile::MMAS_M };
+    enum { MMAS_N = Mma_tile::MMAS_N };
+
+    // The number of groups of warp such that we have at most 4 warps writing consecutive elements.
+    enum { GROUPS = fmha::Div_up<Cta_tile::WARPS_N, 4>::VALUE };
+    // The number of elements that we are going to store per row.
+    enum { ELEMENTS_PER_ROW = Cta_tile::WARPS_N / GROUPS };
+    // The number of rows.
+    enum { ROWS = Cta_tile::M * GROUPS };
+    // The total number of elements.
+    enum { ELEMENTS = ROWS * ELEMENTS_PER_ROW };
+
+    // Ctor.
+    template<typename Params>
+    inline __device__ Softmax_base(const Params &params, void *smem, int bidb, int tidx)
+        :  // packed_mask_ptr_(reinterpret_cast<const char*>(params.packed_mask_ptr)),
+          smem_(reinterpret_cast<float *>(smem)), tidx_(tidx) {
+
+        // Move to the 1st mask loaded by the thread+ tidx;
+        // packed_mask_ptr_ += bidb * params.packed_mask_stride_in_bytes + tidx * sizeof(uint32_t);
+
+        // Extract the position in the warp.
+        int warp = tidx / Cta_tile::THREADS_PER_WARP;
+        int lane = tidx % Cta_tile::THREADS_PER_WARP;
+
+        // Decompose the warp index into M and N.
+        int warp_m = warp % Cta_tile::WARPS_M;
+        int warp_n = warp / Cta_tile::WARPS_M;
+
+        // Decompose the warp-n index into group/position-inside-the-group.
+        int warp_g = warp_n / ELEMENTS_PER_ROW;
+        int warp_i = warp_n % ELEMENTS_PER_ROW;
+
+        // The location written by the threads.
+        int write_row = warp_g * (ROWS / GROUPS) + warp_m * Mma_tile::M_PER_MMA + lane / 4;
+        int write_col = warp_i;
+
+        // Assemble the write pointer.
+        smem_write_ = &smem_[write_row * ELEMENTS_PER_ROW + write_col];
+
+        // Assemble the read pointer.
+        smem_read_ = &smem_[warp_m * Mma_tile::M_PER_MMA + lane / 4];
+    }
+
+    template<typename Mask>
+    inline __device__ void apply_mask(const Mask &mask) {
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M; ++mi ) {
+            #pragma unroll
+            for( int ii = 0; ii < 2; ++ii ) {
+                #pragma unroll
+                for( int ni = 0; ni < MMAS_N; ++ni ) {
+                    #pragma unroll
+                    for( int jj = 0; jj < 4; ++jj ) {
+                        if( !mask.is_valid(mi, ni, ii, jj) ) {
+                            elt_[2 * mi + ii][4 * ni + jj] = -INFINITY;
+                        }
+                    }
+                }
+            }
+        }
+    }
+
+    // Apply the exp to all the elements.
+    inline __device__ void apply_exp(const float (&max)[MMAS_M * 2]) {
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M * 2; ++mi ) {
+            #pragma unroll
+            for( int ni = 0; ni < MMAS_N * 4; ++ni ) {
+                elt_[mi][ni] = apply_exp_(elt_[mi][ni], max[mi]);
+            }
+        }
+    }
+
+    // Do a CTA-wide reduction.
+    template<typename Functor>
+    inline __device__ void reduce_1x4(float (&dst)[MMAS_M * 2]) {
+
+#if defined(USE_SAME_SUM_ORDER_IN_SOFTMAX_AS_REF_CODE)
+        if( Functor::IS_SUM ) {
+            // Apply the summation inside the thread.
+            float tmp[MMAS_M * 2][2];
+            #pragma unroll
+            for( int mi = 0; mi < MMAS_M * 2; ++mi ) {
+                tmp[mi][0] = 0.f;
+                tmp[mi][1] = 0.f;
+                #pragma unroll
+                for( int ni = 0; ni < MMAS_N; ++ni ) {
+                    tmp[mi][0] += elt_[mi][4 * ni + 0];
+                    tmp[mi][0] += elt_[mi][4 * ni + 1];
+                    tmp[mi][1] += elt_[mi][4 * ni + 2];
+                    tmp[mi][1] += elt_[mi][4 * ni + 3];
+                }
+                dst[mi] = tmp[mi][0] + tmp[mi][1];
+            }
+        } else
+#endif  // defined(USE_SAME_SUM_ORDER_IN_SOFTMAX_AS_REF_CODE)
+        {
+            // Apply the functor for each row inside a thread.
+            #pragma unroll
+            for( int mi = 0; mi < MMAS_M * 2; ++mi ) {
+                dst[mi] = elt_[mi][0];
+                #pragma unroll
+                for( int ni = 1; ni < MMAS_N * 4; ++ni ) {
+                    dst[mi] = Functor::apply(dst[mi], elt_[mi][ni]);
+                }
+            }
+        }
+
+        // Apply the functor for each row inside each group of 4 threads.
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M * 2; ++mi ) {
+            dst[mi] = Functor::apply(dst[mi], __shfl_xor_sync(uint32_t(-1), dst[mi], 1));
+            __syncwarp();
+            dst[mi] = Functor::apply(dst[mi], __shfl_xor_sync(uint32_t(-1), dst[mi], 2));
+            __syncwarp();
+        }
+
+        // Store the different values.
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M; ++mi ) {
+            if( tidx_ % 4 == 0 ) {
+                smem_write_[(mi * Mma_tile::M_PER_MMA_PER_CTA + 0) * ELEMENTS_PER_ROW] = dst[2 * mi + 0];
+                smem_write_[(mi * Mma_tile::M_PER_MMA_PER_CTA + 8) * ELEMENTS_PER_ROW] = dst[2 * mi + 1];
+            }
+        }
+
+        // Make sure the values are in shared memory.
+        __syncthreads();
+
+        // Load 8 values (one for each warp). The /8 corresponds to /(4*2) where 4 is from the
+        // float4.
+        float4 tmp[1];
+        if( tidx_ < Cta_tile::M ) {
+            tmp[0] = reinterpret_cast<const float4 *>(&smem_[0 * ELEMENTS / 2])[tidx_];
+        }
+
+        // Compute the reduction of those 8 values in a binary-tree fashion.
+        tmp[0].x = Functor::apply(tmp[0].x, tmp[0].y);
+        tmp[0].z = Functor::apply(tmp[0].z, tmp[0].w);
+        tmp[0].x = Functor::apply(tmp[0].x, tmp[0].z);
+
+        // Make sure we can write to shared memory.
+        __syncthreads();
+
+        // Store the value back to shared memory.
+        if( tidx_ < Cta_tile::M ) {
+            smem_[tidx_] = tmp[0].x;
+        }
+
+        // Make sure the data is in shared memory.
+        __syncthreads();
+
+        // Finally read the values.
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M; ++mi ) {
+            dst[2 * mi + 0] = smem_read_[mi * Mma_tile::M_PER_MMA_PER_CTA + 0];
+            dst[2 * mi + 1] = smem_read_[mi * Mma_tile::M_PER_MMA_PER_CTA + 8];
+        }
+    }
+
+    // Do a CTA-wide reduction.
+    template<typename Functor>
+    inline __device__ void reduce_1x8(float (&dst)[MMAS_M * 2]) {
+
+#if defined(USE_SAME_SUM_ORDER_IN_SOFTMAX_AS_REF_CODE)
+        if( Functor::IS_SUM ) {
+            // Apply the summation inside the thread.
+            float tmp[MMAS_M * 2][2];
+            #pragma unroll
+            for( int mi = 0; mi < MMAS_M * 2; ++mi ) {
+                tmp[mi][0] = 0.f;
+                tmp[mi][1] = 0.f;
+                #pragma unroll
+                for( int ni = 0; ni < MMAS_N; ++ni ) {
+                    tmp[mi][0] += elt_[mi][4 * ni + 0];
+                    tmp[mi][0] += elt_[mi][4 * ni + 1];
+                    tmp[mi][1] += elt_[mi][4 * ni + 2];
+                    tmp[mi][1] += elt_[mi][4 * ni + 3];
+                }
+                dst[mi] = tmp[mi][0] + tmp[mi][1];
+            }
+        } else
+#endif  // defined(USE_SAME_SUM_ORDER_IN_SOFTMAX_AS_REF_CODE)
+        {
+            // Apply the functor for each row inside a thread.
+            #pragma unroll
+            for( int mi = 0; mi < MMAS_M * 2; ++mi ) {
+                dst[mi] = elt_[mi][0];
+                #pragma unroll
+                for( int ni = 1; ni < MMAS_N * 4; ++ni ) {
+                    dst[mi] = Functor::apply(dst[mi], elt_[mi][ni]);
+                }
+            }
+        }
+
+        // Apply the functor for each row inside each group of 4 threads.
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M * 2; ++mi ) {
+            dst[mi] = Functor::apply(dst[mi], __shfl_xor_sync(uint32_t(-1), dst[mi], 1));
+            __syncwarp();
+            dst[mi] = Functor::apply(dst[mi], __shfl_xor_sync(uint32_t(-1), dst[mi], 2));
+            __syncwarp();
+        }
+
+        // Store the different values.
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M; ++mi ) {
+            if( tidx_ % 4 == 0 ) {
+                smem_write_[(mi * Mma_tile::M_PER_MMA_PER_CTA + 0) * ELEMENTS_PER_ROW] = dst[2 * mi + 0];
+                smem_write_[(mi * Mma_tile::M_PER_MMA_PER_CTA + 8) * ELEMENTS_PER_ROW] = dst[2 * mi + 1];
+            }
+        }
+
+        // Make sure the values are in shared memory.
+        __syncthreads();
+
+        // Load 8 values (one for each warp). The /8 corresponds to /(4*2) where 4 is from the
+        // float4.
+        float4 tmp[2];
+        if( tidx_ < Cta_tile::M ) {
+            tmp[0] = reinterpret_cast<const float4 *>(&smem_[0 * ELEMENTS / 2])[tidx_];
+            tmp[1] = reinterpret_cast<const float4 *>(&smem_[1 * ELEMENTS / 2])[tidx_];
+        }
+
+        // Compute the reduction of those 8 values in a binary-tree fashion.
+        tmp[0].x = Functor::apply(tmp[0].x, tmp[0].y);
+        tmp[0].z = Functor::apply(tmp[0].z, tmp[0].w);
+        tmp[1].x = Functor::apply(tmp[1].x, tmp[1].y);
+        tmp[1].z = Functor::apply(tmp[1].z, tmp[1].w);
+        tmp[0].x = Functor::apply(tmp[0].x, tmp[0].z);
+        tmp[1].x = Functor::apply(tmp[1].x, tmp[1].z);
+        tmp[0].x = Functor::apply(tmp[0].x, tmp[1].x);
+
+        // Make sure we can write to shared memory.
+        __syncthreads();
+
+        // Store the value back to shared memory.
+        if( tidx_ < Cta_tile::M ) {
+            smem_[tidx_] = tmp[0].x;
+        }
+
+        // Make sure the data is in shared memory.
+        __syncthreads();
+
+        // Finally read the values.
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M; ++mi ) {
+            dst[2 * mi + 0] = smem_read_[mi * Mma_tile::M_PER_MMA_PER_CTA + 0];
+            dst[2 * mi + 1] = smem_read_[mi * Mma_tile::M_PER_MMA_PER_CTA + 8];
+        }
+    }
+
+    // Do a CTA-wide reduction.
+    template<typename Functor>
+    inline __device__ void reduce(float (&dst)[MMAS_M * 2]) {
+        static_assert(Cta_tile::WARPS_M == 1 && (Cta_tile::WARPS_N == 4 || Cta_tile::WARPS_N == 8));
+        if( Cta_tile::WARPS_M == 1 && Cta_tile::WARPS_N == 4 ) {
+            reduce_1x4<Functor>(dst);
+        } else if( Cta_tile::WARPS_M == 1 && Cta_tile::WARPS_N == 8 ) {
+            reduce_1x8<Functor>(dst);
+        } else {
+            assert(false);
+        }
+
+        // Make sure we are done reading from shared memory.
+        __syncthreads();
+    }
+
+    // Scale all the elements.
+    inline __device__ void scale(const float (&sum)[MMAS_M * 2]) {
+        // Precompute the inverse sum to normalize. Without -use_fast_math, it makes a huge deal.
+        float inv_sum[MMAS_M * 2];
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M * 2; ++mi ) {
+            inv_sum[mi] = (sum[mi] == 0.f || sum[mi] != sum[mi]) ? 1.f : 1.f / sum[mi];
+        }
+
+        // Update the values.
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M * 2; ++mi ) {
+            #pragma unroll
+            for( int ni = 0; ni < MMAS_N * 4; ++ni ) {
+                elt_[mi][ni] *= inv_sum[mi];
+            }
+        }
+    }
+
+    // The pointer to the mask.
+    const char *packed_mask_ptr_;
+    // Shared memory for the CTA-wide reduction.
+    float *smem_, *smem_write_, *smem_read_;
+    // The current thread index.
+    int tidx_;
+    // The elements.
+    float elt_[MMAS_M * 2][MMAS_N * 4];
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Cta_tile, typename Kernel_traits>
+struct Softmax : public Softmax_base<Cta_tile, Kernel_traits> {
+
+    // The base class.
+    using Base = Softmax_base<Cta_tile, Kernel_traits>;
+    // The fragment.
+    using Fragment_a = fmha::Fragment_a<fmha::Row>;
+
+    static_assert(Fragment_a::NUM_REGS == 4);
+
+    // The MMAs.
+    enum { MMAS_M = Base::MMAS_M };
+    enum { MMAS_N = Base::MMAS_N };
+
+    // The accumulators.
+    using Accumulator = fmha::Fragment_accumulator;
+    using Accumulator_out = Fragment<uint16_t, 8>;
+    static_assert(Accumulator_out::NUM_REGS == 4);
+
+    static_assert(std::is_same<Accumulator::Data_type, float>::value);
+
+    // Ctor.
+    template<typename Params>
+    inline __device__ Softmax(const Params &params, void *smem, int bidb, int tidx)
+        : Base(params, smem, bidb, tidx), params_scale_bmm1_(params.scale_bmm1) {
+    }
+
+    // Store the tile after softmax.
+    template<typename Gmem_tile>
+    inline __device__ void store(Gmem_tile &gmem_tile) {
+        Accumulator_out acc[MMAS_M][MMAS_N];
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M; ++mi ) {
+            #pragma unroll
+            for( int ni = 0; ni < MMAS_N; ++ni ) {
+
+                // The elements.
+                float tmp_00 = this->elt_[2 * mi + 0][4 * ni + 0];
+                float tmp_01 = this->elt_[2 * mi + 0][4 * ni + 1];
+                float tmp_02 = this->elt_[2 * mi + 0][4 * ni + 2];
+                float tmp_03 = this->elt_[2 * mi + 0][4 * ni + 3];
+                float tmp_10 = this->elt_[2 * mi + 1][4 * ni + 0];
+                float tmp_11 = this->elt_[2 * mi + 1][4 * ni + 1];
+                float tmp_12 = this->elt_[2 * mi + 1][4 * ni + 2];
+                float tmp_13 = this->elt_[2 * mi + 1][4 * ni + 3];
+
+                // Transform to accumulators.
+                acc[mi][ni].reg(0) = fmha::float2_to_half2(tmp_00, tmp_01);
+                acc[mi][ni].reg(1) = fmha::float2_to_half2(tmp_10, tmp_11);
+                acc[mi][ni].reg(2) = fmha::float2_to_half2(tmp_02, tmp_03);
+                acc[mi][ni].reg(3) = fmha::float2_to_half2(tmp_12, tmp_13);
+            }
+        }
+
+        // Delegate to the gmem tile to store.
+        gmem_tile.store(acc);
+    }
+
+    // Pack the data to a fragment for the next GEMM.
+    template<int K, int M>
+    inline __device__ void pack(Fragment_a (&dst)[K][M]) const {
+        #pragma unroll
+        for( int mi = 0; mi < M; ++mi ) {
+            #pragma unroll
+            for( int ki = 0; ki < K; ++ki ) {
+
+                // 1st row - 4 elements per row.
+                float tmp_00 = this->elt_[2 * mi + 0][4 * ki + 0];
+                float tmp_01 = this->elt_[2 * mi + 0][4 * ki + 1];
+                float tmp_02 = this->elt_[2 * mi + 0][4 * ki + 2];
+                float tmp_03 = this->elt_[2 * mi + 0][4 * ki + 3];
+
+                // 2nd row - 4 elements per row.
+                float tmp_10 = this->elt_[2 * mi + 1][4 * ki + 0];
+                float tmp_11 = this->elt_[2 * mi + 1][4 * ki + 1];
+                float tmp_12 = this->elt_[2 * mi + 1][4 * ki + 2];
+                float tmp_13 = this->elt_[2 * mi + 1][4 * ki + 3];
+
+                // Pack to 4 registers.
+                dst[ki][mi].reg(0) = fmha::float2_to_half2(tmp_00, tmp_01);
+                dst[ki][mi].reg(1) = fmha::float2_to_half2(tmp_10, tmp_11);
+                dst[ki][mi].reg(2) = fmha::float2_to_half2(tmp_02, tmp_03);
+                dst[ki][mi].reg(3) = fmha::float2_to_half2(tmp_12, tmp_13);
+            }
+        }
+    }
+
+    // Scale FP32 fragments
+    inline __device__ void unpack(const Accumulator (&acc)[MMAS_M][MMAS_N]) {
+        const float scalef = reinterpret_cast<const float &>(this->params_scale_bmm1_);
+
+        #pragma unroll
+        for( int mi = 0; mi < MMAS_M; ++mi ) {
+            #pragma unroll
+            for( int ni = 0; ni < MMAS_N; ++ni ) {
+                // 1st row - 4 elements per row.
+                this->elt_[2 * mi + 0][4 * ni + 0] = acc[mi][ni].elt(0) * scalef;
+                this->elt_[2 * mi + 0][4 * ni + 1] = acc[mi][ni].elt(1) * scalef;
+                this->elt_[2 * mi + 0][4 * ni + 2] = acc[mi][ni].elt(4) * scalef;
+                this->elt_[2 * mi + 0][4 * ni + 3] = acc[mi][ni].elt(5) * scalef;
+                // 2nd row - 4 elements per row.
+                this->elt_[2 * mi + 1][4 * ni + 0] = acc[mi][ni].elt(2) * scalef;
+                this->elt_[2 * mi + 1][4 * ni + 1] = acc[mi][ni].elt(3) * scalef;
+                this->elt_[2 * mi + 1][4 * ni + 2] = acc[mi][ni].elt(6) * scalef;
+                this->elt_[2 * mi + 1][4 * ni + 3] = acc[mi][ni].elt(7) * scalef;
+            }
+        }
+    }
+    const uint32_t params_scale_bmm1_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/utils.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/utils.h
new file mode 100644
index 0000000000..2bd03e9b36
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha/utils.h
@@ -0,0 +1,953 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include <assert.h>
+#include <stdint.h>
+#include <stdlib.h>
+
+extern "C" __device__ uint32_t __nvvm_get_smem_pointer(void *ptr);
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct Row {};  
+struct Col {};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int M, bool = (M & (M-1)) == 0 >
+struct Next_power_of_two {
+};
+
+template< int M >
+struct Next_power_of_two<  M, true > { enum { VALUE =   M }; };
+template<>
+struct Next_power_of_two<  3, false> { enum { VALUE =   4 }; };
+template<>
+struct Next_power_of_two<  5, false> { enum { VALUE =   8 }; };
+template<>
+struct Next_power_of_two<  6, false> { enum { VALUE =   8 }; };
+template<>
+struct Next_power_of_two<  7, false> { enum { VALUE =   8 }; };
+template<>
+struct Next_power_of_two<  9, false> { enum { VALUE =  16 }; };
+template<>
+struct Next_power_of_two< 10, false> { enum { VALUE =  16 }; };
+template<>
+struct Next_power_of_two< 11, false> { enum { VALUE =  16 }; };
+template<>
+struct Next_power_of_two< 12, false> { enum { VALUE =  16 }; };
+template<>
+struct Next_power_of_two< 13, false> { enum { VALUE =  16 }; };
+template<>
+struct Next_power_of_two< 14, false> { enum { VALUE =  16 }; };
+template<>
+struct Next_power_of_two< 15, false> { enum { VALUE =  16 }; };
+template<>
+struct Next_power_of_two< 24, false> { enum { VALUE =  32 }; };
+template<>
+struct Next_power_of_two< 48, false> { enum { VALUE =  64 }; };
+template<>
+struct Next_power_of_two< 80, false> { enum { VALUE = 128 }; };
+template<>
+struct Next_power_of_two< 96, false> { enum { VALUE = 128 }; };
+template<>
+struct Next_power_of_two<112, false> { enum { VALUE = 128 }; };
+template<>
+struct Next_power_of_two<144, false> { enum { VALUE = 256 }; };
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N, bool = (N & (N-1)) == 0 >
+struct Prev_power_of_two {
+};
+
+template< int N >
+struct Prev_power_of_two< N, true > { enum { VALUE = N }; };
+template<>
+struct Prev_power_of_two< 3, false> { enum { VALUE = 2 }; };
+template<>
+struct Prev_power_of_two< 5, false> { enum { VALUE = 4 }; };
+template<>
+struct Prev_power_of_two< 6, false> { enum { VALUE = 4 }; };
+template<>
+struct Prev_power_of_two< 7, false> { enum { VALUE = 4 }; };
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int M, int N >
+struct Div_up {
+    enum { VALUE = (M + N-1) / N };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int A, int B >
+struct Max {
+    enum { VALUE = A >= B ? A : B };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int A, int B, int C >
+struct Max_3 {
+    enum { VALUE = Max<Max<A, B>::VALUE, C>::VALUE };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int A, int B >
+struct Min {
+    enum { VALUE = A <= B ? A : B };
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int SIZE_IN_BYTES >
+struct Uint_from_size_in_bytes {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+struct Uint_from_size_in_bytes<1> {
+    using Type = uint8_t;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+struct Uint_from_size_in_bytes<2> {
+    using Type = uint16_t;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+struct Uint_from_size_in_bytes<4> {
+    using Type = uint32_t;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+struct Uint_from_size_in_bytes<8> {
+    using Type = uint2;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+struct Uint_from_size_in_bytes<16> {
+    using Type = uint4;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int WARPS_M, int WARPS_N, int WARPS_K >
+struct Warp_masks {
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+struct Warp_masks<8, 1, 1> { enum { M = 0xe0, N = 0x00, K = 0x00 }; };
+template<>
+struct Warp_masks<4, 2, 1> { enum { M = 0x60, N = 0x80, K = 0x00 }; };
+template<>
+struct Warp_masks<4, 1, 2> { enum { M = 0x60, N = 0x00, K = 0x80 }; };
+template<>
+struct Warp_masks<4, 1, 1> { enum { M = 0x60, N = 0x00, K = 0x00 }; };
+template<>
+struct Warp_masks<2, 4, 1> { enum { M = 0x20, N = 0xc0, K = 0x00 }; };
+template<>
+struct Warp_masks<2, 2, 2> { enum { M = 0x20, N = 0x40, K = 0x80 }; };
+template<>
+struct Warp_masks<2, 2, 1> { enum { M = 0x20, N = 0x40, K = 0x00 }; };
+template<>
+struct Warp_masks<2, 1, 2> { enum { M = 0x20, N = 0x00, K = 0x40 }; };
+template<>
+struct Warp_masks<2, 1, 1> { enum { M = 0x20, N = 0x00, K = 0x00 }; };
+template<>
+struct Warp_masks<1, 8, 1> { enum { M = 0x00, N = 0xe0, K = 0x00 }; };
+template<>
+struct Warp_masks<1, 4, 2> { enum { M = 0x00, N = 0x60, K = 0x80 }; };
+template<>
+struct Warp_masks<1, 4, 1> { enum { M = 0x00, N = 0x60, K = 0x00 }; };
+template<>
+struct Warp_masks<1, 2, 2> { enum { M = 0x00, N = 0x20, K = 0x40 }; };
+template<>
+struct Warp_masks<1, 2, 1> { enum { M = 0x00, N = 0x20, K = 0x00 }; };
+template<>
+struct Warp_masks<1, 1, 4> { enum { M = 0x00, N = 0x00, K = 0x60 }; };
+template<>
+struct Warp_masks<1, 1, 2> { enum { M = 0x00, N = 0x00, K = 0x20 }; };
+template<>
+struct Warp_masks<1, 1, 1> { enum { M = 0x00, N = 0x00, K = 0x00 }; };
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename T >
+inline __device__ __host__ T div_up(T m, T n) {
+    return (m + n-1) / n;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline int clz(int x) {
+    for( int i = 31; i >= 0; --i ) {
+        if( (1 << i) & x ) {
+            return 31 - i;
+        }
+    }
+    return 32;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline int find_log_2(int x, bool round_up = false) {
+    int a = 31 - clz(x);
+    if( round_up ) {
+        a += (x & (x-1)) ? 1 : 0;
+    }
+    return a;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t hadd2(uint32_t a, uint32_t b) {
+    uint32_t c;
+    asm volatile("add.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b));
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t hmin2(uint32_t a, uint32_t b) {
+    uint32_t c;
+    asm volatile("min.f16x2 %0, %1, %2;" : "=r"(c) : "r"(a), "r"(b));
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t hmul2(uint32_t a, uint32_t b) {
+    uint32_t c;
+    asm volatile("mul.f16x2 %0, %1, %2;\n" : "=r"(c) : "r"(a), "r"(b));
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint2 hmul4(uint2 a, uint2 b) {
+    uint2 c;
+    c.x = hmul2(a.x, b.x);
+    c.y = hmul2(a.y, b.y);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint4 hmul8(uint4 a, uint4 b) {
+    uint4 c;
+    c.x = hmul2(a.x, b.x);
+    c.y = hmul2(a.y, b.y);
+    c.z = hmul2(a.z, b.z);
+    c.w = hmul2(a.w, b.w);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint4 hmul8(uint32_t a, uint4 b) {
+    uint4 c;
+    c.x = hmul2(a, b.x);
+    c.y = hmul2(a, b.y);
+    c.z = hmul2(a, b.z);
+    c.w = hmul2(a, b.w);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t hrelu2(uint32_t x, uint32_t lb = 0) {
+    uint32_t res;
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+    asm volatile( "max.f16x2 %0, %1, %2;\n" : "=r"(res) : "r"(x), "r"(lb));
+#else
+    const uint32_t zero = 0u;
+    asm volatile( \
+        "{\n" \
+        "\t .reg .f16x2 sela;\n" \
+        "\t set.gtu.u32.f16x2 sela, %1, %2;\n" \
+        "\t and.b32 %0, sela, %1;\n" 
+        "}\n" : "=r"(res) : "r"(x), "r"(zero));
+#endif
+    return res;
+}
+static inline __device__ uint32_t habs2(uint32_t x) {
+    uint32_t res;
+    asm volatile( "abs.f16x2 %0, %1;\n" : "=r"(res) : "r"(x));
+    return res;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+template< typename T >
+static inline __device__ T clamp(T x, T lb, T ub) {
+    return x < lb ? lb : (x > ub ? ub : x);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint16_t clamp_to_zero(uint16_t x) {
+    uint16_t mask;
+    asm volatile("set.gtu %0, %1, 0;" : "=h"(mask) : "h"(x));
+    return mask & x;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint16_t float_to_half(float f) {
+    uint16_t h;
+    asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(h) : "f"(f));
+    return h;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t float2_to_half2(float a, float b) {
+    uint32_t c;
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+    asm volatile("cvt.rn.f16x2.f32 %0, %1, %2;\n" : "=r"(c) : "f"(b), "f"(a));
+#else
+    uint16_t lo = float_to_half(a);
+    uint16_t hi = float_to_half(b);
+    asm volatile("mov.b32 %0, {%1, %2};\n" : "=r"(c) : "h"(lo), "h"(hi));
+#endif
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t float_to_half2(float a) {
+    return float2_to_half2(a,a);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t float2_to_half2(const float2 &f) {
+    return float2_to_half2(f.x, f.y);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint2 float4_to_half4(float x, float y, float z, float w) {
+    uint2 d;
+    d.x = float2_to_half2(x, y);
+    d.y = float2_to_half2(z, w);
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t hfma2(uint32_t a, uint32_t b, uint32_t c) {
+    uint32_t d;
+    asm volatile("fma.rn.f16x2 %0, %1, %2, %3;\n" : "=r"(d) : "r"(a), "r"(b), "r"(c));
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t hfma2_relu(uint32_t a, uint32_t b, uint32_t c) {
+    uint32_t d;
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 800
+    asm volatile("fma.rn.f16x2.relu %0, %1, %2, %3;" : "=r"(d) : "r"(a), "r"(b), "r"(c));
+#else
+    d = hrelu2(hfma2(a, b, c));
+#endif
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t h0_h0(uint32_t x) {
+    uint32_t y;
+    asm volatile("{.reg .f16 lo, hi; mov.b32 {lo, hi}, %1; mov.b32 %0, {lo, lo};}\n" 
+        : "=r"(y) : "r"(x)); 
+    return y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ float h0_to_float(uint32_t h2) {
+    float f;
+    asm volatile("{\n" \
+        ".reg .f16 lo, hi;\n" \
+        "mov.b32 {lo, hi}, %1;\n" \
+        "cvt.f32.f16 %0, lo;\n" \
+        "}\n" : "=f"(f) : "r"(h2));
+    return f;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t h1_h1(uint32_t x) {
+    uint32_t y;
+    asm volatile("{.reg .f16 lo, hi; mov.b32 {lo, hi}, %1; mov.b32 %0, {hi, hi};}\n" 
+        : "=r"(y) : "r"(x)); 
+    return y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint16_t hadd(uint16_t a, uint16_t b) {
+    uint16_t d;
+    asm volatile("add.f16 %0, %1, %2;" : "=h"(d) : "h"(a), "h"(b));
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint32_t hadd(uint32_t a, uint32_t b) {
+    return hadd2(a, b);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint2 hadd4(uint2 a, uint2 b) {
+    uint2 c;
+    c.x = hadd2(a.x, b.x);
+    c.y = hadd2(a.y, b.y);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint2 hadd(uint2 a, uint2 b) {
+    return hadd4(a, b);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint4 hadd8(uint4 a, uint4 b) {
+    uint4 c;
+    c.x = hadd2(a.x, b.x);
+    c.y = hadd2(a.y, b.y);
+    c.z = hadd2(a.z, b.z);
+    c.w = hadd2(a.w, b.w);
+    return c;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint4 fadd4(uint4 a, uint4 b) {
+    float4 c;
+    c.x = reinterpret_cast<const float&>(a.x) + reinterpret_cast<const float&>(b.x);
+    c.y = reinterpret_cast<const float&>(a.y) + reinterpret_cast<const float&>(b.y);
+    c.z = reinterpret_cast<const float&>(a.z) + reinterpret_cast<const float&>(b.z);
+    c.w = reinterpret_cast<const float&>(a.w) + reinterpret_cast<const float&>(b.w);
+    return reinterpret_cast<const uint4&>(c);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint4 hadd(uint4 a, uint4 b) {
+    return hadd8(a, b);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ float half_to_float(uint16_t h) {
+    float f;
+    asm volatile("cvt.f32.f16 %0, %1;\n" : "=f"(f) : "h"(h));
+    return f;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ float2 half2_to_float2(uint32_t x) {
+    uint16_t lo, hi;
+    asm volatile("mov.b32 {%0, %1}, %2;\n" : "=h"(lo), "=h"(hi) : "r"(x));
+    return make_float2(half_to_float(lo), half_to_float(hi));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ void half2_to_float2(float &x, float &y, uint32_t h) {
+    float2 tmp = half2_to_float2(h);
+    x = tmp.x;
+    y = tmp.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint16_t hfma(uint16_t a, uint16_t b, uint16_t c) {
+    uint16_t d;
+    asm volatile("fma.rn.f16 %0, %1, %2, %3;" : "=h"(d) : "h"(a), "h"(b), "h"(c));
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ uint16_t hmul(uint16_t a, uint16_t b) {
+    uint16_t d;
+    asm volatile("mul.f16 %0, %1, %2;" : "=h"(d) : "h"(a), "h"(b));
+    return d;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline __device__ float sigmoid(float x) {
+    return 1.f / (1.f + expf(-x));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void clear(uint16_t &dst) {
+    dst = uint16_t(0);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void clear(uint32_t &dst) {
+    dst = 0u;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void clear(uint2 &dst) {
+    dst = make_uint2(0u, 0u);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void clear(uint4 &dst) {
+    dst = make_uint4(0u, 0u, 0u, 0u);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// P R E D I C A T E   P A C K I N G
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+enum { BYTES_PER_REG = 4, PREDS_PER_BYTE = 4, PREDS_PER_REG = BYTES_PER_REG * PREDS_PER_BYTE };
+
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// G E N E R I C   P R E D I C A T E D   L D G S T S
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N, int M, typename Functor >
+inline __device__ void load_(Functor &fct, const uint32_t (&preds)[M]) {
+
+    // The number of complete bytes (where we use all the predicates in a byte).
+    enum { COMPLETE = N / PREDS_PER_BYTE };
+    // Make sure we did allocate enough predicates.
+    static_assert(Div_up<COMPLETE, BYTES_PER_REG>::VALUE <= M, "");
+    // The remainder.
+    enum { REMAINDER = N - COMPLETE * PREDS_PER_BYTE };
+    // Make sure we got the math right and the remainder is between 0 and 3.
+    static_assert(REMAINDER >= 0 && REMAINDER <= 3, "");
+    // The mask to extract the predicates.
+    enum { COMPLETE_MASK = (1 << PREDS_PER_BYTE) - 1 };
+
+    // Clear the fetch registers.
+    #pragma unroll
+    for( int ii = 0; ii < N; ++ii ) {
+        fct.clear(ii);
+    }
+
+    // Run complete steps.
+    bool p[PREDS_PER_BYTE];
+    #pragma unroll
+    for( int ii = 0; ii < COMPLETE; ++ii ) {
+
+        // The predicate.
+        uint32_t reg = preds[ii / BYTES_PER_REG];
+
+        // Extract the predicates.
+        #pragma unroll
+        for( int jj = 0; jj < PREDS_PER_BYTE; ++jj ) {
+            uint32_t mask = 1u << (ii % BYTES_PER_REG * 8 + jj);
+            p[jj] = (reg & mask) != 0u;
+        }
+
+        // Issue the loads.
+        #pragma unroll
+        for( int jj = 0; jj < PREDS_PER_BYTE; ++jj ) {
+            fct.load(ii * PREDS_PER_BYTE + jj, p[jj]);
+        }
+    }
+
+    // Skip the rest of the code if we do not have a remainder.
+    if( REMAINDER > 0 ) {
+
+        // The mask to extract the predicates.
+        enum { REMAINDER_MASK = (1 << REMAINDER) - 1 };
+
+        // The predicate register.
+        uint32_t reg = preds[COMPLETE / BYTES_PER_REG];
+
+        // Extract the predicates.
+        #pragma unroll
+        for( int jj = 0; jj < PREDS_PER_BYTE; ++jj ) {
+            uint32_t mask = 1u << (COMPLETE % BYTES_PER_REG * 8 + jj);
+            p[jj] = (reg & mask) != 0u;
+        }
+
+        // Issue the loads.
+        #pragma unroll
+        for( int ii = 0; ii < REMAINDER; ++ii ) {
+            fct.load(COMPLETE * PREDS_PER_BYTE + ii, p[ii]);
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int M, typename Functor >
+inline __device__ void load_(Functor &fct, uint32_t preds) {
+    uint32_t tmp[1] = { preds };
+    load_<M>(fct, tmp);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// L D G
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldg(uint8_t &dst, const void *ptr) {
+    dst = *reinterpret_cast<const uint8_t*>(ptr);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldg(uint16_t &dst, const void *ptr) {
+    dst = *reinterpret_cast<const uint16_t*>(ptr);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldg(uint32_t &dst, const void *ptr) {
+    dst = *reinterpret_cast<const uint32_t*>(ptr);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldg(uint2 &dst, const void *ptr) {
+    dst = *reinterpret_cast<const uint2*>(ptr);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldg(uint4 &dst, const void *ptr) {
+    dst = *reinterpret_cast<const uint4*>(ptr);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Data_type, int N >
+struct Ldg_functor {
+    // Ctor.
+    inline __device__ Ldg_functor(Data_type (&fetch)[N], const void* (&ptrs)[N])
+        : fetch_(fetch), ptrs_(ptrs) {
+    }
+
+    // Clear the element.
+    inline __device__ void clear(int ii) {
+        fmha::clear(fetch_[ii]);
+    }
+
+    // Trigger the loads.
+    inline __device__ void load(int ii, bool p) {
+        if( p ) {
+            ldg(fetch_[ii], ptrs_[ii]);
+        }
+    }
+
+    // The fetch registers.
+    Data_type (&fetch_)[N];
+    // The pointers.
+    const void* (&ptrs_)[N];
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Data_type, int N, int M >
+inline __device__ void ldg_(Data_type (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {
+    Ldg_functor<Data_type, N> fct(fetch, ptrs);
+    load_<N>(fct, preds);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N, int M >
+inline __device__ void ldg(uint8_t (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {
+    ldg_<uint8_t, N>(fetch, ptrs, preds);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N, int M >
+inline __device__ void ldg(uint16_t (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {
+    ldg_<uint16_t, N>(fetch, ptrs, preds);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N, int M >
+inline __device__ void ldg(uint32_t (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {
+    ldg_<uint32_t, N>(fetch, ptrs, preds);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N, int M >
+inline __device__ void ldg(uint2 (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {
+    ldg_<uint2, N>(fetch, ptrs, preds);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N, int M >
+inline __device__ void ldg(uint4 (&fetch)[N], const void* (&ptrs)[N], uint32_t (&preds)[M]) {
+    ldg_<uint4, N>(fetch, ptrs, preds);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// L D S
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void lds(uint16_t &dst, uint32_t ptr) {
+    asm volatile("ld.shared.b16 %0, [%1];\n" : "=h"(dst) : "r"(ptr));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void lds(uint32_t &dst, uint32_t ptr) {
+    asm volatile("ld.shared.b32 %0, [%1];\n" : "=r"(dst) : "r"(ptr));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void lds(uint2 &dst, uint32_t ptr) {
+    asm volatile("ld.shared.v2.b32 {%0, %1}, [%2];\n" : "=r"(dst.x), "=r"(dst.y) : "r"(ptr));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void lds(uint4 &dst, uint32_t ptr) {
+    asm volatile("ld.shared.v4.b32 {%0, %1, %2, %3}, [%4];\n"
+        : "=r"(dst.x)
+        , "=r"(dst.y)
+        , "=r"(dst.z)
+        , "=r"(dst.w)
+        :  "r"(ptr));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// L D S M
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldsm(uint32_t &dst, uint32_t ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730
+    asm volatile("ldmatrix.sync.aligned.m8n8.x1.shared.b16 {%0}, [%1];\n"
+        : "=r"(dst) : "r"(ptr));
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldsmt(uint32_t &dst, uint32_t ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730
+    asm volatile("ldmatrix.sync.aligned.m8n8.x1.trans.shared.b16 {%0}, [%1];\n"
+        : "=r"(dst) : "r"(ptr));
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldsm(uint2 &dst, uint32_t ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730
+    asm volatile("ldmatrix.sync.aligned.m8n8.x2.shared.b16 {%0, %1}, [%2];\n"
+        : "=r"(dst.x), "=r"(dst.y) : "r"(ptr));
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldsmt(uint2 &dst, uint32_t ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730
+    asm volatile("ldmatrix.sync.aligned.m8n8.x2.trans.shared.b16 {%0, %1}, [%2];\n"
+        : "=r"(dst.x), "=r"(dst.y) : "r"(ptr));
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldsm(uint4 &dst, uint32_t ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+        : "=r"(dst.x), "=r"(dst.y), "=r"(dst.z), "=r"(dst.w) : "r"(ptr));
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void ldsmt(uint4 &dst, uint32_t ptr) {
+#if defined(__CUDA_ARCH__) && __CUDA_ARCH__ >= 730
+    asm volatile("ldmatrix.sync.aligned.m8n8.x4.trans.shared.b16 {%0, %1, %2, %3}, [%4];\n"
+        : "=r"(dst.x), "=r"(dst.y), "=r"(dst.z), "=r"(dst.w) : "r"(ptr));
+#endif
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// S T G
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void stg(void *ptr, uint8_t val) {
+    *reinterpret_cast<uint8_t*>(ptr) = val;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void stg(void *ptr, uint16_t val) {
+    *reinterpret_cast<uint16_t*>(ptr) = val;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void stg(void *ptr, uint32_t val) {
+    *reinterpret_cast<uint32_t*>(ptr) = val;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void stg(void *ptr, uint2 val) {
+    *reinterpret_cast<uint2*>(ptr) = val;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void stg(void *ptr, uint4 val) {
+    *reinterpret_cast<uint4*>(ptr) = val;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+//
+// S T S
+//
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void sts(uint32_t ptr, uint16_t val) {
+    asm volatile("st.shared.b16 [%0], %1;\n" : : "r"(ptr), "h"(val));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void sts(uint32_t ptr, uint32_t val) {
+    asm volatile("st.shared.b32 [%0], %1;\n" : : "r"(ptr), "r"(val));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void sts(uint32_t ptr, uint2 val) {
+    asm volatile("st.shared.v2.b32 [%0], {%1, %2};\n"
+        :
+        : "r"(ptr)
+        , "r"(val.x)
+        , "r"(val.y));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void sts(uint32_t ptr, uint4 val) {
+    asm volatile("st.shared.v4.b32 [%0], {%1, %2, %3, %4};\n"
+        :
+        : "r"(ptr)
+        , "r"(val.x)
+        , "r"(val.y)
+        , "r"(val.z)
+        , "r"(val.w));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename Data_type, int N >
+inline __device__ void sts_(uint32_t (&ptrs)[N], const Data_type (&data)[N]) {
+    #pragma unroll
+    for( int ii = 0; ii < N; ++ii ) {
+        sts(ptrs[ii], data[ii]);
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+inline __device__ void sts(uint32_t (&ptrs)[N], const uint16_t (&data)[N]) {
+    sts_<uint16_t, N>(ptrs, data);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+inline __device__ void sts(uint32_t (&ptrs)[N], const uint32_t (&data)[N]) {
+    sts_<uint32_t, N>(ptrs, data);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+inline __device__ void sts(uint32_t (&ptrs)[N], const uint2 (&data)[N]) {
+    sts_<uint2, N>(ptrs, data);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+inline __device__ void sts(uint32_t (&ptrs)[N], const uint4 (&data)[N]) {
+    sts_<uint4, N>(ptrs, data);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_128_64_kernel.sm80.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_128_64_kernel.sm80.cu
new file mode 100644
index 0000000000..bed79b8483
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_128_64_kernel.sm80.cu
@@ -0,0 +1,60 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include "fmha.h"
+#include "fmha_dgrad_kernel_1xN_reload.h"
+
+using Kernel_traits = FMHA_kernel_traits< 128, 64, 16, 1, 4, 0x08u>;
+
+extern "C" __global__ void fmha_dgrad_fp16_128_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::compute_dv_1xN<Kernel_traits>(params);
+    fmha::compute_dq_dk_1xN<Kernel_traits>(params);
+}
+
+void run_fmha_dgrad_fp16_128_64_sm80(const Fused_multihead_attention_fprop_params &params, cudaStream_t stream) {
+
+    constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);
+    constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;
+    constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;
+    constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;
+
+    using Smem_tile_s = fmha::Smem_tile_mma_transposed< Kernel_traits::Cta_tile_p>;
+    constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;
+    static_assert(smem_size_s == 16 * 128 * 2);
+    static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);
+
+    constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;
+    constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;
+    constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);
+
+    if( smem_size >= 48 * 1024 ) {
+        FMHA_CHECK_CUDA(cudaFuncSetAttribute(
+            fmha_dgrad_fp16_128_64_sm80_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+    dim3 grid(params.h, params.b);
+    fmha_dgrad_fp16_128_64_sm80_kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_256_64_kernel.sm80.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_256_64_kernel.sm80.cu
new file mode 100644
index 0000000000..f93c628cad
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_256_64_kernel.sm80.cu
@@ -0,0 +1,60 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include "fmha.h"
+#include "fmha_dgrad_kernel_1xN_reload.h"
+
+using Kernel_traits = FMHA_kernel_traits< 256, 64, 16, 1, 4, 0x08u>;
+
+extern "C" __global__ void fmha_dgrad_fp16_256_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::compute_dv_1xN<Kernel_traits>(params);
+    fmha::compute_dq_dk_1xN<Kernel_traits>(params);
+}
+
+void run_fmha_dgrad_fp16_256_64_sm80(const Fused_multihead_attention_fprop_params &params, cudaStream_t stream) {
+
+    constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);
+    constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;
+    constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;
+    constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;
+
+    using Smem_tile_s = fmha::Smem_tile_mma_transposed< Kernel_traits::Cta_tile_p>;
+    constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;
+    static_assert(smem_size_s == 16 * 256 * 2);
+    static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);
+
+    constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;
+    constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;
+    constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);
+
+    if( smem_size >= 48 * 1024 ) {
+        FMHA_CHECK_CUDA(cudaFuncSetAttribute(
+            fmha_dgrad_fp16_256_64_sm80_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+    dim3 grid(params.h, params.b);
+    fmha_dgrad_fp16_256_64_sm80_kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_384_64_kernel.sm80.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_384_64_kernel.sm80.cu
new file mode 100644
index 0000000000..ff117f06db
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_384_64_kernel.sm80.cu
@@ -0,0 +1,60 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include "fmha.h"
+#include "fmha_dgrad_kernel_1xN_reload.h"
+
+using Kernel_traits = FMHA_kernel_traits< 384, 64, 16, 1, 8, 0x08u>;
+
+extern "C" __global__ void fmha_dgrad_fp16_384_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::compute_dv_1xN<Kernel_traits>(params);
+    fmha::compute_dq_dk_1xN<Kernel_traits>(params);
+}
+
+void run_fmha_dgrad_fp16_384_64_sm80(const Fused_multihead_attention_fprop_params &params, cudaStream_t stream) {
+
+    constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);
+    constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;
+    constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;
+    constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;
+
+    using Smem_tile_s = fmha::Smem_tile_mma_transposed< Kernel_traits::Cta_tile_p>;
+    constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;
+    static_assert(smem_size_s == 16 * 384 * 2);
+    static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);
+
+    constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;
+    constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;
+    constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);
+
+    if( smem_size >= 48 * 1024 ) {
+        FMHA_CHECK_CUDA(cudaFuncSetAttribute(
+            fmha_dgrad_fp16_384_64_sm80_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+    dim3 grid(params.h, params.b);
+    fmha_dgrad_fp16_384_64_sm80_kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_512_64_kernel.sm80.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_512_64_kernel.sm80.cu
new file mode 100644
index 0000000000..8fbc44d308
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_512_64_kernel.sm80.cu
@@ -0,0 +1,105 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include "fmha.h"
+#include "fmha_dgrad_kernel_1xN_reload.h"
+#include "fmha_dgrad_kernel_1xN_reload_nl.h"
+
+using Kernel_traits = FMHA_kernel_traits< 512, 64, 16, 1, 8, 0x08u>;
+
+extern "C" __global__ void fmha_dgrad_fp16_512_64_sm80_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::compute_dv_1xN<Kernel_traits>(params);
+    fmha::compute_dq_dk_1xN<Kernel_traits>(params);
+}
+
+template<int CHUNKS>
+__global__
+void fmha_dgrad_fp16_512_64_sm80_nl_kernel(Fused_multihead_attention_fprop_params params){
+    fmha::compute_dv_1xN_nl<CHUNKS, Kernel_traits>(params);
+    fmha::compute_dq_dk_1xN_nl<CHUNKS, Kernel_traits>(params);
+}
+
+void run_fmha_dgrad_fp16_512_64_sm80(const Fused_multihead_attention_fprop_params &params, cudaStream_t stream) {
+
+    constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);
+    constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;
+    constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;
+    constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;
+
+    using Smem_tile_s = fmha::Smem_tile_mma_transposed< Kernel_traits::Cta_tile_p>;
+    constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;
+    static_assert(smem_size_s == 16 * 512 * 2);
+    static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);
+
+    constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;
+    constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;
+    constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);
+
+    if( smem_size >= 48 * 1024 ) {
+        FMHA_CHECK_CUDA(cudaFuncSetAttribute(
+            fmha_dgrad_fp16_512_64_sm80_kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+    dim3 grid(params.h, params.b);
+    fmha_dgrad_fp16_512_64_sm80_kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);
+}
+
+void run_fmha_dgrad_fp16_512_64_sm80_nl(const Fused_multihead_attention_fprop_params &params, const int num_chunks, cudaStream_t stream) {
+
+    constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);
+    constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;
+    constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;
+    constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;
+
+    using Smem_tile_s = fmha::Smem_tile_mma_transposed<Kernel_traits::Cta_tile_p>;
+    constexpr int smem_size_s = Smem_tile_s::BYTES_PER_TILE;
+    static_assert(smem_size_s == 16 * 512 * 2);
+    static_assert(smem_size_o == 16 * 64 * 4 * Kernel_traits::Cta_tile_p::WARPS_N);
+
+    constexpr int smem_size_dv = smem_size_s + 2 * smem_size_q + smem_size_v + smem_size_softmax;
+    constexpr int smem_size_dq_dk = smem_size_s + smem_size_o + smem_size_q + smem_size_v;
+    constexpr int smem_size = std::max(smem_size_dv, smem_size_dq_dk);
+
+    auto kernel = fmha_dgrad_fp16_512_64_sm80_nl_kernel<2>;
+       
+    if( num_chunks == 2 ) {
+        kernel = fmha_dgrad_fp16_512_64_sm80_nl_kernel<2>;
+    }else if( num_chunks == 3 ) {
+        kernel = fmha_dgrad_fp16_512_64_sm80_nl_kernel<3>;
+    } else {
+        assert(false && "Unsupperted number of chunks");
+    }
+
+    if( smem_size >= 48 * 1024 ) {
+        FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+
+    dim3 grid(params.h, params.b, num_chunks);
+
+    kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);
+
+    FMHA_CHECK_CUDA(cudaPeekAtLastError());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_kernel_1xN_reload.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_kernel_1xN_reload.h
new file mode 100644
index 0000000000..ddcd6210e8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_kernel_1xN_reload.h
@@ -0,0 +1,558 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include "fmha_kernel.h"
+#include <fmha/kernel_traits.h>
+#include <fmha/gemm.h>
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Kernel_traits, typename Params>
+inline __device__ void compute_dv_1xN(const Params &params) {
+
+    // The description of the CTA tile for the 1st batched GEMM.
+    using Cta_tile_p = typename Kernel_traits::Cta_tile_p;
+    // The description of the CTA tile for the 2nd batched GEMM.
+    using Cta_tile_dv =
+        fmha::Cta_tile_extd<Cta_tile_p::N, Cta_tile_p::K, Cta_tile_p::M, Cta_tile_p::WARPS_N, 1, Cta_tile_p::WARPS_M>;
+
+    static_assert(Cta_tile_dv::M == 512 || Cta_tile_dv::M == 384 || Cta_tile_dv::M == 256 || Cta_tile_dv::M == 128);
+    static_assert(Cta_tile_dv::N == 64);
+    static_assert(Cta_tile_dv::K == 16);
+
+    // The MMA tile for the 1st GEMM.
+    using Mma_tile_p = fmha::Hmma_tile<Cta_tile_p>;
+    // The MMA tile for the 2nd GEMM.
+    using Mma_tile_dv = fmha::Hmma_tile<Cta_tile_dv>;
+
+    // The global memory tile to load Q.
+    using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;
+    // The shared memory tile to swizzle Q.
+    // using Smem_tile_q = typename Kernel_traits::Smem_tile_q;
+    using Smem_tile_q = fmha::Smem_tile_a<Cta_tile_p, fmha::Row, Gmem_tile_q::BYTES_PER_LDG, 2>;
+    // The shared memory tile to reload Q as fragment b.
+    using Smem_tile_qt = fmha::Smem_tile_b<Cta_tile_dv, fmha::Row, Gmem_tile_q::BYTES_PER_LDG, 2>;
+
+    // The global memory tile to load K.
+    using Gmem_tile_k = typename Kernel_traits::Gmem_tile_k;
+    // The shared memory tile to swizzle K.
+    using Smem_tile_k = typename Kernel_traits::Smem_tile_k;
+
+    // The global memory tile to load V.
+    using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle V.
+    using Smem_tile_v = typename Kernel_traits::Smem_tile_v;
+
+    // The global memory tile to store O.
+    using Gmem_tile_o = typename Kernel_traits::Gmem_tile_o;
+    // The shared memory tile to swizzle O.
+    using Smem_tile_o = typename Kernel_traits::Smem_tile_o;
+
+    // The global memory tile to store dV.
+    using Gmem_tile_dv = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle dV.
+    using Smem_tile_dv = fmha::Smem_tile_mma_epilogue<Cta_tile_dv>;
+    static_assert(Smem_tile_dv::NUM_LDS == Gmem_tile_dv::LDGS);
+    static_assert(Smem_tile_dv::THREADS_PER_ROW == Gmem_tile_dv::THREADS_PER_ROW);
+
+    using Gmem_tile_s = typename Kernel_traits::Gmem_tile_s;
+    using Smem_tile_st = typename Kernel_traits::Smem_tile_st;
+    using Gmem_tile_do = typename Kernel_traits::Gmem_tile_do;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+
+    // The block index for the batch.
+    const int bidb = blockIdx.y;
+    // The block index for the head.
+    const int bidh = blockIdx.x;
+    // The thread index.
+    const int tidx = threadIdx.x;
+
+    const BlockInfoPadded<Kernel_traits::THREADS> binfo(params, bidb, bidh, tidx);
+    if( binfo.stop_early() )
+        return;
+    Mask<Cta_tile_p> mask(params, binfo, tidx);
+
+    // Allocate the global memory tile loader for Q.
+    Gmem_tile_do gmem_q(params, binfo, tidx);  // treating dout as Q
+    // Allocate the shared memory tile loader for Q.
+    Smem_tile_q smem_q(&smem_[0], tidx);
+    Smem_tile_qt smem_qt(&smem_[0], tidx);
+    Smem_tile_st smem_s(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE], tidx);
+
+    // Allocate the global memory tile loader for K.
+    Gmem_tile_k gmem_k(params, 2, binfo, tidx);  // treating V as K
+    // Allocate the shared memory tile loader for K.
+    Smem_tile_k smem_k(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);
+
+    // Trigger the loads for Q.
+    gmem_q.load(smem_q);
+    // Trigger the loads for K.
+    gmem_k.load(smem_k);
+
+    // Commit the data for Q and K to shared memory.
+    gmem_q.commit(smem_q);
+    gmem_k.commit(smem_k);
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Load the fragments for Q.
+    typename Smem_tile_q::Fragment frag_q[2][Mma_tile_p::MMAS_M];
+    smem_q.load(frag_q[0], 0);
+
+    typename Smem_tile_qt::Fragment frag_qt[2][Mma_tile_dv::MMAS_N];
+    static_assert(Smem_tile_qt::Fragment::NUM_REGS == 4);
+    static_assert(Mma_tile_dv::MMAS_K == 1);
+    smem_qt.load(frag_qt[0], 0);
+
+    // Load the fragments for K. We keep the data in registers during the entire kernel.
+    typename Smem_tile_k::Fragment frag_k[2][Mma_tile_p::MMAS_N];
+    smem_k.load(frag_k[0], 0);
+
+    enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };
+
+    Gmem_tile_s gmem_s(params.s_ptr, params, tidx);
+
+    // Create the object to do the softmax.
+    using Softmax = fmha::Softmax<Cta_tile_p, Kernel_traits>;
+    Softmax softmax(
+        params, &smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_st::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE], bidb, tidx);
+
+    enum { THREADS_PER_ROW = 32 };
+    enum { M = Mma_tile_p::MMAS_M };
+    enum { N = Mma_tile_p::MMAS_N };
+
+    // Declare the accumulators for the 2nd gemm.
+    fmha::Fragment_accumulator acc_dv[Mma_tile_dv::MMAS_M][Mma_tile_dv::MMAS_N];
+    fmha::Clear_accumulator<fmha::Accumulator_type, Cta_tile_dv::WARPS_K>::apply(acc_dv);
+
+    enum { STEPS = Cta_tile_p::N / Cta_tile_p::M };
+    // Load over the entire sequence length.
+    for( int l = 0; l < STEPS; l++ ) {
+        const int loop = l * Cta_tile_p::M;
+        if( loop >= binfo.actual_seqlen )
+            break;
+
+        // Load S
+        uint4 s_regs[M][N];
+        gmem_s.load(s_regs, mask);
+        fmha::Fragment_accumulator acc_p[Mma_tile_p::MMAS_M][Mma_tile_p::MMAS_N];
+        fmha::Clear_accumulator<fmha::Accumulator_type, Cta_tile_p::WARPS_K>::apply(acc_p);
+        // Do this part of P^T = (Q * K^T)^T.
+        #pragma unroll
+        for( int ki = 1; ki < Mma_tile_p::MMAS_K; ++ki ) {
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_q.load(frag_q[ki & 1], ki);
+            smem_k.load(frag_k[ki & 1], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1) & 1]);
+        }
+
+        // Store s * dmask to smem for transpose
+        smem_s.store(s_regs);
+
+        // Declare the accumulators for the 1st gemm.
+        // Do the final stage of math.
+        {
+            int ki = Mma_tile_p::MMAS_K;
+            fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1) & 1]);
+        }
+        // Trigger the load for the next Q values. We're using double buffering, so reading qt is safe
+        if( l < STEPS - 1) {
+            smem_q.move_to_next_write_buffer();
+            gmem_q.move();
+            gmem_q.load(smem_q);
+        }
+
+
+        // Convert from the accumulator type to FP32 for Softmax.
+        softmax.unpack(acc_p);
+
+        float s_mat[2 * M][4 * N];
+
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ni = 0; ni < N; ni++ ) {
+                uint4 &dst = s_regs[mi][ni];
+                fmha::half2_to_float2(s_mat[2 * mi + 0][4 * ni + 0], s_mat[2 * mi + 0][4 * ni + 1], dst.x);
+                fmha::half2_to_float2(s_mat[2 * mi + 0][4 * ni + 2], s_mat[2 * mi + 0][4 * ni + 3], dst.y);
+                fmha::half2_to_float2(s_mat[2 * mi + 1][4 * ni + 0], s_mat[2 * mi + 1][4 * ni + 1], dst.z);
+                fmha::half2_to_float2(s_mat[2 * mi + 1][4 * ni + 2], s_mat[2 * mi + 1][4 * ni + 3], dst.w);
+            }
+        }
+
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ii = 0; ii < 2; ii++ ) {
+                #pragma unroll
+                for( int ni = 0; ni < N; ni++ ) {
+                    #pragma unroll
+                    for( int jj = 0; jj < 4; jj++ ) {
+                        float & s_dmask = s_mat[2 * mi + ii][4 * ni + jj];
+                        const bool drop = reinterpret_cast<const uint32_t &>(s_dmask) & 0x80000000;
+                        const float d_s = drop ? 0.f : softmax.elt_[2 * mi + ii][4 * ni + jj] * params.rp_dropout;
+                        s_dmask = fabsf(s_dmask);
+                        softmax.elt_[2 * mi + ii][4 * ni + jj] = d_s * fabsf(s_dmask);
+                    }
+                }
+            }
+        }
+
+        float p_sum[2 * M];
+        softmax.template reduce<fmha::Sum_>(p_sum);
+
+        const float scalef = reinterpret_cast<const float &>(params.scale_softmax);
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ii = 0; ii < 2; ii++ ) {
+                #pragma unroll
+                for( int ni = 0; ni < N; ni++ ) {
+                    #pragma unroll
+                    for( int jj = 0; jj < 4; jj++ ) {
+                        softmax.elt_[2 * mi + ii][4 * ni + jj] -= p_sum[2 * mi + ii] * (s_mat[2 * mi + ii][4 * ni + jj]) ;
+                        softmax.elt_[2 * mi + ii][4 * ni + jj] *= scalef;
+                    }
+                }
+            }
+        }
+        typename Smem_tile_st::Fragment frag_s[Mma_tile_dv::MMAS_K][Mma_tile_dv::MMAS_M];
+        smem_s.load(frag_s);
+        for( int ki = 0; ki < Mma_tile_dv::MMAS_K; ki++ ) {
+            for( int mi = 0; mi < Mma_tile_dv::MMAS_M; mi++ ) {
+                for( int ii = 0; ii < Smem_tile_st::Fragment::NUM_REGS; ii++ ) {
+                    frag_s[ki][mi].reg(ii) = fmha::hmul2(frag_s[ki][mi].reg(ii), params.scale_dropout);
+                    frag_s[ki][mi].reg(ii) = fmha::hrelu2(frag_s[ki][mi].reg(ii));
+                }
+            }
+        }
+
+        gmem_s.store(softmax.elt_, mask);
+        gmem_s.move();
+
+        #pragma unroll
+        for( int ki = 1; ki < Mma_tile_dv::MMAS_K; ++ki ) {
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_qt.load(frag_qt[ki & 1], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_dv, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);
+        }
+
+        // Do the final stage of math.
+        {
+            int ki = Mma_tile_dv::MMAS_K;
+            fmha::gemm(acc_dv, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);
+        }
+        // Commit the values for Q into shared memory.
+        if(l < STEPS - 1) {
+            gmem_q.commit(smem_q);
+        }
+
+        // Make sure we are reading from the correct buffer.
+        smem_q.move_to_next_read_buffer();
+        smem_qt.move_to_next_read_buffer();
+
+        // Make sure the data is in shared memory.
+        __syncthreads();
+
+        // Trigger the loads for the values of Q for the next iteration.
+        smem_q.load(frag_q[0], 0);
+        smem_k.load(frag_k[0], 0);
+        smem_qt.load(frag_qt[0], 0);
+
+    }  // Outer loop over the sequence length.
+
+    // Epilogue swizzle for dV
+    Smem_tile_dv smem_dv(&smem_[Kernel_traits::Smem_tile_q::BYTES_PER_TILE], tidx);
+    smem_dv.store(acc_dv);
+
+    __syncthreads();
+    uint4 dv_out[Smem_tile_dv::NUM_LDS];
+    smem_dv.load(dv_out);
+    Qkv_params dv_params;
+    dv_params.qkv_ptr = params.dqkv_ptr;
+    dv_params.qkv_stride_in_bytes = params.qkv_stride_in_bytes;
+    dv_params.h = params.h;
+    Gmem_tile_dv gmem_dv(dv_params, 2, binfo, tidx);
+    gmem_dv.store(dv_out);
+}
+
+template<typename Kernel_traits, typename Params>
+inline __device__ void compute_dq_dk_1xN(const Params &params) {
+
+    // The description of the CTA tile for the 1st batched GEMM.
+    using Cta_tile_p = typename Kernel_traits::Cta_tile_p;
+    using Cta_tile_o = typename Kernel_traits::Cta_tile_o;
+    // The description of the CTA tile for the 2nd batched GEMM.
+    using Cta_tile_dk =
+        fmha::Cta_tile_extd<Cta_tile_p::N, Cta_tile_p::K, Cta_tile_p::M, Cta_tile_p::WARPS_N, 1, Cta_tile_p::WARPS_M>;
+    static_assert(Cta_tile_dk::M == 512 || Cta_tile_dk::M == 384 || Cta_tile_dk::M == 256 || Cta_tile_dk::M == 128);
+    static_assert(Cta_tile_dk::N == 64);
+    static_assert(Cta_tile_dk::K == 16);
+
+    // The MMA tile for the 1st GEMM.
+    using Mma_tile_p = fmha::Hmma_tile<Cta_tile_p>;
+    using Mma_tile_o = fmha::Hmma_tile<Cta_tile_o>;
+    // The MMA tile for the 2nd GEMM.
+    using Mma_tile_dk = fmha::Hmma_tile<Cta_tile_dk>;
+
+    // The global memory tile to load Q.
+    using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;
+    // The shared memory tile to swizzle Q.
+    using Smem_tile_q = typename Kernel_traits::Smem_tile_q;
+
+    // The global memory tile to load K.
+    using Gmem_tile_k = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle K.
+    using Smem_tile_k = typename Kernel_traits::Smem_tile_v;  // K is used like V in fprop
+
+    // The global memory tile to load V.
+    using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle V.
+    using Smem_tile_v = typename Kernel_traits::Smem_tile_v;
+
+    // The global memory tile to store O.
+    // using Gmem_tile_o = typename Kernel_traits::Gmem_tile_o;
+    using Gmem_tile_o = fmha::Gmem_tile_dq<Cta_tile_o>;
+    // The shared memory tile to swizzle O.
+    using Smem_tile_o = typename Kernel_traits::Smem_tile_o;
+
+    // The global memory tile to store dK.
+    using Gmem_tile_dk = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle dK.
+    using Smem_tile_dk = fmha::Smem_tile_mma_epilogue<Cta_tile_dk>;
+    static_assert(Smem_tile_dk::NUM_LDS == Gmem_tile_dk::LDGS);
+    static_assert(Smem_tile_dk::THREADS_PER_ROW == Gmem_tile_dk::THREADS_PER_ROW);
+
+    // The shared memory tile to reload Q transposed.
+    using Smem_tile_qt = fmha::Smem_tile_b<Cta_tile_dk, fmha::Row, Gmem_tile_q::BYTES_PER_LDG, 1>;
+
+    using Gmem_tile_s = typename Kernel_traits::Gmem_tile_s;
+
+    using Smem_tile_st = typename Kernel_traits::Smem_tile_st;
+
+
+    enum { M = Mma_tile_p::MMAS_M };
+    enum { N = Mma_tile_p::MMAS_N };
+    static_assert(M == Mma_tile_o::MMAS_M);
+    static_assert(N == Mma_tile_o::MMAS_K);
+    // Shared memory.
+    extern __shared__ char smem_[];
+
+    // The block index for the batch.
+    const int bidb = blockIdx.y;
+    // The block index for the head.
+    const int bidh = blockIdx.x;
+    // The thread index.
+    const int tidx = threadIdx.x;
+
+    const BlockInfoPadded<Kernel_traits::THREADS> binfo(params, bidb, bidh, tidx);
+    if( binfo.stop_early() )
+        return;
+
+    Mask<Cta_tile_p> mask(params, binfo, tidx);
+
+    // Allocate the global memory tile loader for Q.
+    Gmem_tile_q gmem_q(params, 0, binfo, tidx);
+    // Allocate the shared memory tile loader for Q.
+    Smem_tile_q smem_q(&smem_[0], tidx);
+    Smem_tile_qt smem_qt(&smem_[0], tidx);
+    Smem_tile_st smem_s(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE], tidx);
+
+    // Allocate the global memory tile loader for K.
+    Gmem_tile_k gmem_k(params, 1, binfo, tidx);
+    // Allocate the shared memory tile loader for K.
+    Smem_tile_k smem_k(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);
+
+    // Allocate the global memory tile loader for O.
+    Gmem_tile_o gmem_o(params, binfo, tidx);
+    // Allocate the shared memory tile loader for O. We use the same as K so be careful!!!
+    Smem_tile_o smem_o(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE], tidx);
+
+    // Trigger the loads for Q.
+    gmem_q.load(smem_q);
+    // Trigger the loads for K.
+    gmem_k.load(smem_k);
+
+    Gmem_tile_s gmem_s(params.s_ptr, params, tidx);
+    // Load dP
+    uint4 s_regs[M][N];
+    gmem_s.load(s_regs, mask);
+    gmem_s.move();
+
+    // Commit the data for Q and K to shared memory.
+    gmem_q.commit(smem_q);
+    gmem_k.commit(smem_k);
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    typename Smem_tile_qt::Fragment frag_qt[2][Mma_tile_dk::MMAS_N];
+    smem_qt.load(frag_qt[0], 0);
+    typename Smem_tile_k::Fragment frag_k[2][Mma_tile_o::MMAS_N];
+    smem_k.load(frag_k[0], 0);
+
+    enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };
+
+    enum { THREADS_PER_ROW = 32 };
+    enum { STEPS = Cta_tile_p::N / Cta_tile_p::M };
+
+    // Declare the accumulators for the 2nd gemm.
+    fmha::Fragment_accumulator acc_dk[Mma_tile_dk::MMAS_M][Mma_tile_dk::MMAS_N];
+    fmha::Clear_accumulator<fmha::Accumulator_type, Cta_tile_dk::WARPS_K>::apply(acc_dk);
+
+    // Load over the entire sequence length.
+    for( int l=0;l<STEPS;l++) {
+        const int loop = l * Cta_tile_p::M;
+        if( loop >= binfo.actual_seqlen )
+            break;
+
+        // Pack dP as Fragment_a
+        fmha::Fragment_a<fmha::Row> frag_p[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_M];
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ni = 0; ni < N; ni++ ) {
+                uint4 &dst = s_regs[mi][ni];
+                frag_p[ni][mi].reg(0) = dst.x;  // row 0, cols 0,1
+                frag_p[ni][mi].reg(1) = dst.z;  // row 8, cols 0,1
+                frag_p[ni][mi].reg(2) = dst.y;  // row 0, cols 8,9
+                frag_p[ni][mi].reg(3) = dst.w;  // row 8, cols 8,9
+            }
+        }
+
+        // Declare the accumulators for the 1st gemm.
+        fmha::Fragment_accumulator acc_o[Mma_tile_o::MMAS_M][Mma_tile_o::MMAS_N];
+        fmha::Clear_accumulator<fmha::Accumulator_type, Cta_tile_o::WARPS_K>::apply(acc_o);
+
+        // Do this part of O = P^T * V^T. dQ = dP x dK
+        #pragma unroll
+        for( int ki = 1; ki < Mma_tile_o::MMAS_K; ++ki ) {
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_k.load(frag_k[ki & 1], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_o, frag_p[ki - 1], frag_k[(ki - 1) & 1]);
+        }
+
+        // Do the final stage of math.
+        {
+            int ki = Mma_tile_o::MMAS_K;
+            fmha::gemm(acc_o, frag_p[ki - 1], frag_k[(ki - 1) & 1]);
+        }
+
+        // Store dP to smem for transpose
+        smem_s.store(s_regs);
+        if(l < STEPS - 1) {
+            // Load next part of S
+            gmem_s.load(s_regs, mask);
+            gmem_s.move();
+            smem_q.move_to_next_write_buffer();
+            gmem_q.move();
+            gmem_q.load(smem_q);
+        }
+        // Loop over MMAS_M.
+        #pragma unroll
+        for( int ii = 0; ii < Gmem_tile_o::LOOPS; ++ii ) {
+
+            // Swizzle the elements and do the final reduction.
+            smem_o.store(acc_o, ii);
+
+            // Make sure the data is in shared memory.
+            __syncthreads();
+
+            // Load from shared memory.
+            uint4 out[Gmem_tile_o::STGS_PER_LOOP];
+            smem_o.load(out);
+
+            // Make sure the data was read from shared memory.
+            if( ii < Gmem_tile_o::LOOPS - 1 ) {
+                __syncthreads();
+            }
+
+            // Output the values.
+            gmem_o.store(out, ii);
+        }
+
+        // Move to the next part of the output.
+        gmem_o.move();
+
+        typename Smem_tile_st::Fragment frag_s[Mma_tile_dk::MMAS_K][Mma_tile_dk::MMAS_M];
+        smem_s.load(frag_s);
+
+        #pragma unroll
+        for( int ki = 1; ki < Mma_tile_dk::MMAS_K; ++ki ) {
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_qt.load(frag_qt[ki & 1], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_dk, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);
+        }
+
+        // Do the final stage of math.
+        {
+            int ki = Mma_tile_dk::MMAS_K;
+            fmha::gemm(acc_dk, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);
+        }
+
+        // Commit the values for Q into shared memory.
+        if( l < STEPS - 1) {
+            gmem_q.commit(smem_q);
+        }
+
+        // Make sure the data is in shared memory.
+        __syncthreads();
+
+        // Trigger the loads for the values of Q for the next iteration.
+        smem_qt.load(frag_qt[0], 0);
+        smem_k.load(frag_k[0], 0);
+
+    }  // Outer loop over the sequence length.
+
+    // Epilogue swizzle for dK
+    Smem_tile_dk smem_dk(&smem_[0], tidx);
+    smem_dk.store(acc_dk);
+    __syncthreads();
+    uint4 dk_out[Smem_tile_dk::NUM_LDS];
+    smem_dk.load(dk_out);
+    Qkv_params dk_params;
+    dk_params.qkv_ptr = params.dqkv_ptr;
+    dk_params.qkv_stride_in_bytes = params.qkv_stride_in_bytes;
+    dk_params.h = params.h;
+    Gmem_tile_dk gmem_dk(dk_params, 1, binfo, tidx);
+    gmem_dk.store(dk_out);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_kernel_1xN_reload_nl.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_kernel_1xN_reload_nl.h
new file mode 100644
index 0000000000..5ecbcea441
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_dgrad_kernel_1xN_reload_nl.h
@@ -0,0 +1,571 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include "fmha_kernel.h"
+#include <fmha/kernel_traits.h>
+#include <fmha/gemm.h>
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int CHUNKS, typename Kernel_traits, typename Params>
+inline __device__ void compute_dv_1xN_nl(const Params &params) {
+
+    // The description of the CTA tile for the 1st batched GEMM.
+    using Cta_tile_p = typename Kernel_traits::Cta_tile_p;
+    // The description of the CTA tile for the 2nd batched GEMM.
+    using Cta_tile_dv = fmha::Cta_tile_extd<Cta_tile_p::N, Cta_tile_p::K, Cta_tile_p::M, Cta_tile_p::WARPS_N, 1, Cta_tile_p::WARPS_M>;
+
+    static_assert(Cta_tile_dv::M == 512 || Cta_tile_dv::M == 384 || Cta_tile_dv::M == 256 || Cta_tile_dv::M == 128);
+    static_assert(Cta_tile_dv::N == 64);
+    static_assert(Cta_tile_dv::K == 16);
+
+    // The MMA tile for the 1st GEMM.
+    using Mma_tile_p = fmha::Hmma_tile<Cta_tile_p>;
+    // The MMA tile for the 2nd GEMM.
+    using Mma_tile_dv = fmha::Hmma_tile<Cta_tile_dv>;
+
+    // The global memory tile to load Q.
+    using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;
+    // The shared memory tile to swizzle Q.
+    using Smem_tile_q = fmha::Smem_tile_a<Cta_tile_p, fmha::Row, Gmem_tile_q::BYTES_PER_LDG, 2>;
+    // The shared memory tile to reload Q as fragment b.
+    using Smem_tile_qt = fmha::Smem_tile_b<Cta_tile_dv, fmha::Row, Gmem_tile_q::BYTES_PER_LDG, 2>;
+
+    // The global memory tile to load K.
+    using Gmem_tile_k = typename Kernel_traits::Gmem_tile_k;
+    // The shared memory tile to swizzle K.
+    using Smem_tile_k = typename Kernel_traits::Smem_tile_k;
+
+    // The global memory tile to load V.
+    using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle V.
+    using Smem_tile_v = typename Kernel_traits::Smem_tile_v;
+
+    // The global memory tile to store dV.
+    using Gmem_tile_dv = fmha::Gmem_tile_qkv<typename Kernel_traits::Cta_tile_o, 
+                                             fmha::BITS_PER_ELEMENT_B, 
+                                             Cta_tile_p::N, //S, 
+                                             Cta_tile_p::K, //D, 
+                                             2*CHUNKS>;
+
+    // The shared memory tile to swizzle dV.
+    using Smem_tile_dv = fmha::Smem_tile_mma_epilogue<Cta_tile_dv>;
+    static_assert(Smem_tile_dv::NUM_LDS == Gmem_tile_dv::LDGS);
+    static_assert(Smem_tile_dv::THREADS_PER_ROW == Gmem_tile_dv::THREADS_PER_ROW);
+
+    using Gmem_tile_s = typename Kernel_traits::Gmem_tile_s;
+    using Smem_tile_st = typename Kernel_traits::Smem_tile_st;
+    using Gmem_tile_do = typename Kernel_traits::Gmem_tile_do;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+
+    // The block index for the chunk.
+    const int bidc = blockIdx.z;
+    // The block index for the batch.
+    const int bidb = blockIdx.y;
+    // The block index for the head.
+    const int bidh = blockIdx.x;
+    // The thread index.
+    const int tidx = threadIdx.x;
+
+    const BlockInfoPadded<Kernel_traits::THREADS> binfo(params, bidb, bidh, tidx);
+    if( binfo.stop_early() )
+        return;
+    fmha::Mask<Cta_tile_p> mask(params, binfo, tidx);
+
+    // Allocate the global memory tile loader for Q.
+    Gmem_tile_do gmem_q(params, binfo, tidx);  // treating dout as Q
+    // Allocate the shared memory tile loader for Q.
+    Smem_tile_q smem_q(&smem_[0], tidx);
+    Smem_tile_qt smem_qt(&smem_[0], tidx);
+    Smem_tile_st smem_s(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE], tidx);
+
+    // Allocate the global memory tile loader for K.
+    Gmem_tile_k gmem_k(params, 2, binfo, tidx);  // treating V as K
+    // Allocate the shared memory tile loader for K.
+    Smem_tile_k smem_k(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);
+
+    Gmem_tile_s gmem_s(params.s_ptr, params, tidx);
+
+    using Noloop = Noloop_traits<CHUNKS, Cta_tile_p>;
+
+    Noloop nl_traits(bidc);
+    nl_traits.move_all(gmem_q, gmem_s);
+
+    // Trigger the loads for Q.
+    gmem_q.load(smem_q);
+    // Trigger the loads for K.
+    gmem_k.load(smem_k);
+
+    // Commit the data for Q and K to shared memory.
+    gmem_q.commit(smem_q);
+    gmem_k.commit(smem_k);
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Load the fragments for Q.
+    typename Smem_tile_q::Fragment frag_q[2][Mma_tile_p::MMAS_M];
+    smem_q.load(frag_q[0], 0);
+
+    typename Smem_tile_qt::Fragment frag_qt[2][Mma_tile_dv::MMAS_N];
+    static_assert(Smem_tile_qt::Fragment::NUM_REGS == 4);
+    static_assert(Mma_tile_dv::MMAS_K == 1);
+    smem_qt.load(frag_qt[0], 0);
+
+    // Load the fragments for K. We keep the data in registers during the entire kernel.
+    typename Smem_tile_k::Fragment frag_k[2][Mma_tile_p::MMAS_N];
+    smem_k.load(frag_k[0], 0);
+
+    enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };
+
+    // Create the object to do the softmax.
+    using Softmax = fmha::Softmax<Cta_tile_p, Kernel_traits>;
+    Softmax softmax(
+        params, &smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_st::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE], bidb, tidx);
+
+    enum { THREADS_PER_ROW = 32 };
+    enum { M = Mma_tile_p::MMAS_M };
+    enum { N = Mma_tile_p::MMAS_N };
+
+    // Declare the accumulators for the 2nd gemm.
+    fmha::Fragment_accumulator acc_dv[Mma_tile_dv::MMAS_M][Mma_tile_dv::MMAS_N];
+    fmha::Clear_accumulator<fmha::Accumulator_type, Cta_tile_dv::WARPS_K>::apply(acc_dv);
+
+    // Load over the entire sequence length.
+    for(int l = 0; l < nl_traits.num_steps_;l++) {
+        const int loop = nl_traits.offset_loop_count(l);
+        if( loop >= binfo.actual_seqlen ) break;
+
+        uint4 s_regs[M][N];
+        gmem_s.load(s_regs, mask);
+        fmha::Fragment_accumulator acc_p[Mma_tile_p::MMAS_M][Mma_tile_p::MMAS_N];
+        fmha::Clear_accumulator<fmha::Accumulator_type, Cta_tile_p::WARPS_K>::apply(acc_p);
+        // Do this part of P^T = (Q * K^T)^T.
+        #pragma unroll
+        for( int ki = 1; ki < Mma_tile_p::MMAS_K; ++ki ) {
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_q.load(frag_q[ki & 1], ki);
+            smem_k.load(frag_k[ki & 1], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1) & 1]);
+        }
+
+        smem_s.store(s_regs);
+
+        // Declare the accumulators for the 1st gemm.
+        // Do the final stage of math.
+        {
+            int ki = Mma_tile_p::MMAS_K;
+            fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1) & 1]);
+        }
+        // Trigger the load for the next Q values. We're using double buffering, so reading qt is safe
+        if(l < nl_traits.num_steps_ - 1) {
+            smem_q.move_to_next_write_buffer();
+            gmem_q.move();
+            gmem_q.load(smem_q);
+        }
+        // Convert from the accumulator type to FP32 for Softmax.
+        softmax.unpack(acc_p);
+
+        float s_mat[2 * M][4 * N];
+
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ni = 0; ni < N; ni++ ) {
+                uint4 &dst = s_regs[mi][ni];
+                fmha::half2_to_float2(s_mat[2 * mi + 0][4 * ni + 0], s_mat[2 * mi + 0][4 * ni + 1], dst.x);
+                fmha::half2_to_float2(s_mat[2 * mi + 0][4 * ni + 2], s_mat[2 * mi + 0][4 * ni + 3], dst.y);
+                fmha::half2_to_float2(s_mat[2 * mi + 1][4 * ni + 0], s_mat[2 * mi + 1][4 * ni + 1], dst.z);
+                fmha::half2_to_float2(s_mat[2 * mi + 1][4 * ni + 2], s_mat[2 * mi + 1][4 * ni + 3], dst.w);
+            }
+        }
+
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ii = 0; ii < 2; ii++ ) {
+                #pragma unroll
+                for( int ni = 0; ni < N; ni++ ) {
+                    #pragma unroll
+                    for( int jj = 0; jj < 4; jj++ ) {
+                         float & s_dmask = s_mat[2 * mi + ii][4 * ni + jj];
+                        const bool drop = reinterpret_cast<const uint32_t &>(s_dmask) & 0x80000000;
+                        const float d_s= drop ? 0.f : softmax.elt_[2 * mi + ii][4 * ni + jj] * params.rp_dropout;
+                        s_dmask = fabsf(s_dmask);
+                        softmax.elt_[2 * mi + ii][4 * ni + jj] = d_s * (s_dmask);
+                    }
+                }
+            }
+        }
+
+        float p_sum[2 * M];
+        softmax.template reduce<fmha::Sum_>(p_sum);
+
+        const float scalef = reinterpret_cast<const float &>(params.scale_softmax);
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ii = 0; ii < 2; ii++ ) {
+                #pragma unroll
+                for( int ni = 0; ni < N; ni++ ) {
+                    #pragma unroll
+                    for( int jj = 0; jj < 4; jj++ ) {
+                        softmax.elt_[2 * mi + ii][4 * ni + jj] -= p_sum[2 * mi + ii] * (s_mat[2 * mi + ii][4 * ni + jj]) ;
+                        softmax.elt_[2 * mi + ii][4 * ni + jj] *= scalef;
+                    }
+                }
+            }
+        }
+
+        typename Smem_tile_st::Fragment frag_s[Mma_tile_dv::MMAS_K][Mma_tile_dv::MMAS_M];
+        smem_s.load(frag_s);
+        for( int ki = 0; ki < Mma_tile_dv::MMAS_K; ki++ ) {
+            for( int mi = 0; mi < Mma_tile_dv::MMAS_M; mi++ ) {
+                for( int ii = 0; ii < Smem_tile_st::Fragment::NUM_REGS; ii++ ) {
+                    frag_s[ki][mi].reg(ii) = fmha::hmul2(frag_s[ki][mi].reg(ii), params.scale_dropout);
+                    frag_s[ki][mi].reg(ii) = fmha::hrelu2(frag_s[ki][mi].reg(ii));
+                }
+            }
+        }
+
+        gmem_s.store(softmax.elt_, mask);
+        gmem_s.move();
+
+        static_assert(Mma_tile_dv::MMAS_K == 1);  // DEBUG
+        #pragma unroll
+        for( int ki = 1; ki < Mma_tile_dv::MMAS_K; ++ki ) {
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_qt.load(frag_qt[ki & 1], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_dv, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);
+        }
+
+        // Do the final stage of math.
+        {
+            int ki = Mma_tile_dv::MMAS_K;
+            fmha::gemm(acc_dv, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);
+        }
+        // Commit the values for Q into shared memory.
+        if(l < nl_traits.num_steps_ - 1) {
+            gmem_q.commit(smem_q);
+        }
+
+        // Make sure we are reading from the correct buffer.
+        smem_q.move_to_next_read_buffer();
+        smem_qt.move_to_next_read_buffer();
+
+        // Make sure the data is in shared memory.
+        __syncthreads();
+
+        // Trigger the loads for the values of Q for the next iteration.
+        smem_q.load(frag_q[0], 0);
+        smem_k.load(frag_k[0], 0);
+        smem_qt.load(frag_qt[0], 0);
+
+    }  // Outer loop over the sequence length.
+
+    // Epilogue for dV = (S * D)' * dout'. We're fully exposed to this!
+
+    // Epilogue swizzle for dV
+    Smem_tile_dv smem_dv(&smem_[Kernel_traits::Smem_tile_q::BYTES_PER_TILE], tidx);
+    smem_dv.store(acc_dv);
+
+    __syncthreads();
+
+    uint4 dv_out[Smem_tile_dv::NUM_LDS];
+    smem_dv.load(dv_out);
+    Qkv_params dv_params;
+    dv_params.qkv_ptr = params.dkv_ptr;
+    dv_params.qkv_stride_in_bytes = params.h * 2 * CHUNKS * params.d * sizeof(half);
+    dv_params.h = params.h;
+    Gmem_tile_dv gmem_dv(dv_params, nl_traits.get_idx_dv(), binfo, tidx);
+    gmem_dv.store(dv_out);
+}
+
+template<int CHUNKS, typename Kernel_traits, typename Params>
+inline __device__ void compute_dq_dk_1xN_nl(const Params &params) {
+
+    // The description of the CTA tile for the 1st batched GEMM.
+    using Cta_tile_p = typename Kernel_traits::Cta_tile_p;
+    using Cta_tile_o = typename Kernel_traits::Cta_tile_o;
+    // The description of the CTA tile for the 2nd batched GEMM.
+    using Cta_tile_dk = fmha::Cta_tile_extd<Cta_tile_p::N, Cta_tile_p::K, Cta_tile_p::M, Cta_tile_p::WARPS_N, 1, Cta_tile_p::WARPS_M>;
+
+    static_assert(Cta_tile_dk::M == 512 || Cta_tile_dk::M == 384 || Cta_tile_dk::M == 256 || Cta_tile_dk::M == 128);
+    static_assert(Cta_tile_dk::N == 64);
+    static_assert(Cta_tile_dk::K == 16);
+
+    // The MMA tile for the 1st GEMM.
+    using Mma_tile_p = fmha::Hmma_tile<Cta_tile_p>;
+    using Mma_tile_o = fmha::Hmma_tile<Cta_tile_o>;
+    // The MMA tile for the 2nd GEMM.
+    using Mma_tile_dk = fmha::Hmma_tile<Cta_tile_dk>;
+
+    // The global memory tile to load Q.
+    using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;
+    // The shared memory tile to swizzle Q.
+    using Smem_tile_q = typename Kernel_traits::Smem_tile_q;
+
+    // The global memory tile to load K.
+    using Gmem_tile_k = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle K.
+    using Smem_tile_k = typename Kernel_traits::Smem_tile_v;  // K is used like V in fprop
+
+    // The global memory tile to load V.
+    using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle V.
+    using Smem_tile_v = typename Kernel_traits::Smem_tile_v;
+
+    // The global memory tile to store O.
+    using Gmem_tile_o = Gmem_tile_dq<Cta_tile_o>;
+    // The shared memory tile to swizzle O.
+    using Smem_tile_o = typename Kernel_traits::Smem_tile_o;
+
+    // The global memory tile to store dK.
+    using Gmem_tile_dk = fmha::Gmem_tile_qkv<typename Kernel_traits::Cta_tile_o, 
+                                             fmha::BITS_PER_ELEMENT_B, 
+                                             Cta_tile_p::N, //S, 
+                                             Cta_tile_p::K, //D, 
+                                             2*CHUNKS>;
+
+    // The shared memory tile to swizzle dK.
+    using Smem_tile_dk = fmha::Smem_tile_mma_epilogue<Cta_tile_dk>;
+    static_assert(Smem_tile_dk::NUM_LDS == Gmem_tile_dk::LDGS);
+    static_assert(Smem_tile_dk::THREADS_PER_ROW == Gmem_tile_dk::THREADS_PER_ROW);
+
+    // The shared memory tile to reload Q transposed.
+    using Smem_tile_qt = fmha::Smem_tile_b<Cta_tile_dk, fmha::Row, Gmem_tile_q::BYTES_PER_LDG, 1>;
+
+    // The global memory tile to load dP, stored in S
+    using Gmem_tile_s = Gmem_tile_mma_s<Cta_tile_p>;
+    // The shared memory tile to transpose dP.
+    using Smem_tile_st = Smem_tile_mma_transposed<Cta_tile_p>;  
+
+    using Noloop = Noloop_traits<CHUNKS, Cta_tile_p>;
+
+    enum { M = Mma_tile_p::MMAS_M };
+    enum { N = Mma_tile_p::MMAS_N };
+    static_assert(M == Mma_tile_o::MMAS_M);
+    static_assert(N == Mma_tile_o::MMAS_K);
+    // Shared memory.
+    extern __shared__ char smem_[];
+
+    const int bidc = blockIdx.z;
+    // The block index for the batch.
+    const int bidb = blockIdx.y;
+    // The block index for the head.
+    const int bidh = blockIdx.x;
+    // The thread index.
+    const int tidx = threadIdx.x;
+
+    const BlockInfoPadded<Kernel_traits::THREADS> binfo(params, bidb, bidh, tidx);
+    if( binfo.stop_early() )
+        return;
+
+    fmha::Mask<Cta_tile_p> mask(params, binfo, tidx);
+
+    // Allocate the global memory tile loader for Q.
+    Gmem_tile_q gmem_q(params, 0, binfo, tidx);
+    // Allocate the shared memory tile loader for Q (as B).
+    Smem_tile_qt smem_qt(&smem_[0], tidx);
+    // Allocate the global memory tile loader for dP.
+    Gmem_tile_s gmem_s(params.s_ptr, params, tidx);
+    // Allocate the shared memory tile loader for dP.
+    Smem_tile_st smem_s(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE], tidx);
+
+    // Allocate the global memory tile loader for K.
+    Gmem_tile_k gmem_k(params, 1, binfo, tidx);
+    // Allocate the shared memory tile loader for K.
+    Smem_tile_k smem_k(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);
+
+    // Allocate the global memory tile loader for O.
+    Gmem_tile_o gmem_o(params, binfo, tidx);
+    // Allocate the shared memory tile loader for O. We use the same as K so be careful!!!
+    Smem_tile_o smem_o(&smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE], tidx);
+
+    Noloop nl_traits(bidc);
+
+    nl_traits.move_all(gmem_q, gmem_o, gmem_s);
+
+    // Trigger the loads for Q.
+    gmem_q.load(smem_qt);
+    // Trigger the loads for K.
+    gmem_k.load(smem_k);
+
+    uint4 s_regs[M][N];
+    gmem_s.load(s_regs, mask);
+
+    // Commit the data for Q and K to shared memory.
+    gmem_q.commit(smem_qt);
+    gmem_k.commit(smem_k);
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    typename Smem_tile_qt::Fragment frag_qt[2][Mma_tile_dk::MMAS_N];
+    smem_qt.load(frag_qt[0], 0);
+    typename Smem_tile_k::Fragment frag_k[2][Mma_tile_o::MMAS_N];
+    smem_k.load(frag_k[0], 0);
+
+    enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };
+
+    enum { THREADS_PER_ROW = 32 };
+
+    // Declare the accumulators for the 2nd gemm.
+    fmha::Fragment_accumulator acc_dk[Mma_tile_dk::MMAS_M][Mma_tile_dk::MMAS_N];
+    fmha::Clear_accumulator<fmha::Accumulator_type, Cta_tile_dk::WARPS_K>::apply(acc_dk);
+
+    // Load over the entire sequence length.
+    for(int l=0;l < nl_traits.num_steps_; l++) {
+
+        // Pack dP as Fragment_a
+        fmha::Fragment_a<fmha::Row> frag_p[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_M];
+        #pragma unroll
+        for( int mi = 0; mi < M; mi++ ) {
+            #pragma unroll
+            for( int ni = 0; ni < N; ni++ ) {
+                uint4 &dst = s_regs[mi][ni];
+                frag_p[ni][mi].reg(0) = dst.x;
+                frag_p[ni][mi].reg(1) = dst.z;
+                frag_p[ni][mi].reg(2) = dst.y;
+                frag_p[ni][mi].reg(3) = dst.w;
+            }
+        }
+        smem_s.store(s_regs);
+        if(l < nl_traits.num_steps_- 1) {
+            // Load next part of S
+            gmem_s.move();
+            gmem_s.load(s_regs, mask);
+            // Trigger the load for the next Q values.
+            smem_qt.move_to_next_write_buffer();
+            gmem_q.move();
+            gmem_q.load(smem_qt);
+        }
+        // Declare the accumulators for the 1st gemm.
+        fmha::Fragment_accumulator acc_o[Mma_tile_o::MMAS_M][Mma_tile_o::MMAS_N];
+        fmha::Clear_accumulator<fmha::Accumulator_type, Cta_tile_o::WARPS_K>::apply(acc_o);
+
+        // Do this part of O = P^T * V^T. dQ = dP x dK
+        #pragma unroll
+        for( int ki = 1; ki < Mma_tile_o::MMAS_K; ++ki ) {
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_k.load(frag_k[ki & 1], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_o, frag_p[ki - 1], frag_k[(ki - 1) & 1]);
+        }
+
+        // Do the final stage of math.
+        {
+            int ki = Mma_tile_o::MMAS_K;
+            fmha::gemm(acc_o, frag_p[ki - 1], frag_k[(ki - 1) & 1]);
+        }
+
+        static_assert(Gmem_tile_o::LOOPS == 1); //DEBUG
+        // Loop over MMAS_M.
+        #pragma unroll
+        for( int ii = 0; ii < Gmem_tile_o::LOOPS; ++ii ) {
+
+            // Swizzle the elements and do the final reduction.
+            smem_o.store(acc_o, ii);
+
+            // Make sure the data is in shared memory.
+            __syncthreads();
+
+            // Load from shared memory.
+            uint4 out[Gmem_tile_o::STGS_PER_LOOP];
+            smem_o.load(out);
+
+            // Make sure the data was read from shared memory.
+            if( ii < Gmem_tile_o::LOOPS - 1 ) {
+                __syncthreads();
+            }
+
+            // Output the values.
+            gmem_o.store(out, ii);
+        }
+
+        // Move to the next part of the output.
+        gmem_o.move();
+
+        typename Smem_tile_st::Fragment frag_s[Mma_tile_dk::MMAS_K][Mma_tile_dk::MMAS_M];
+        smem_s.load(frag_s);
+
+        static_assert(Mma_tile_dk::MMAS_K == 1);  // DEBUG
+
+        #pragma unroll
+        for( int ki = 1; ki < Mma_tile_dk::MMAS_K; ++ki ) {
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_qt.load(frag_qt[ki & 1], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_dk, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);
+        }
+
+        // Do the final stage of math.
+        {
+            int ki = Mma_tile_dk::MMAS_K;
+            fmha::gemm(acc_dk, frag_s[(ki - 1)], frag_qt[(ki - 1) & 1]);
+        }
+
+        // Commit the values for Q into shared memory.
+        if(l < nl_traits.num_steps_- 1) {
+            gmem_q.commit(smem_qt);
+            __syncthreads();
+            // Trigger the loads for the values of Q for the next iteration.
+            smem_qt.load(frag_qt[0], 0);
+            smem_k.load(frag_k[0], 0);
+        }
+
+    }  // Outer loop over the sequence length.
+
+    // Epilogue for dK = dP' * dq. We're fully exposed to this!
+
+    // Epilogue swizzle for dK
+    Smem_tile_dk smem_dk(&smem_[0], tidx);
+    smem_dk.store(acc_dk);
+    
+    __syncthreads();
+    
+    uint4 dk_out[Smem_tile_dk::NUM_LDS];
+    smem_dk.load(dk_out);
+    Qkv_params dk_params;
+    dk_params.qkv_ptr = params.dkv_ptr;
+    dk_params.qkv_stride_in_bytes = params.h * 2 * CHUNKS * params.d * sizeof(half);
+    dk_params.h = params.h;
+    Gmem_tile_dk gmem_dk(dk_params, nl_traits.get_idx_dk(), binfo, tidx);
+    gmem_dk.store(dk_out);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_128_64_kernel.sm80.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_128_64_kernel.sm80.cu
new file mode 100644
index 0000000000..cc325080ca
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_128_64_kernel.sm80.cu
@@ -0,0 +1,58 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include "fmha.h"
+#include "fmha_fprop_kernel_1xN.h"
+
+using Kernel_traits = FMHA_kernel_traits< 128, 64, 16, 1, 4, 0x08u>;
+
+extern "C" __global__ void fmha_fprop_fp16_128_64_sm80_train_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::device_1xN<Kernel_traits, true>(params);
+}
+
+extern "C" __global__ void fmha_fprop_fp16_128_64_sm80_predict_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::device_1xN<Kernel_traits, false>(params);
+}
+
+void run_fmha_fp16_128_64_sm80(const Fused_multihead_attention_fprop_params &params, bool is_training, cudaStream_t stream) {
+
+    auto kernel = is_training ? &fmha_fprop_fp16_128_64_sm80_train_kernel : &fmha_fprop_fp16_128_64_sm80_predict_kernel;
+
+    constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);
+    constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;
+    constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;
+    constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;
+
+    constexpr int smem_size = smem_size_q + std::max(smem_size_v, smem_size_o + smem_size_softmax);
+
+    if( smem_size >= 48 * 1024 ) {
+        FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+
+    dim3 grid(params.h, params.b);
+    kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_256_64_kernel.sm80.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_256_64_kernel.sm80.cu
new file mode 100644
index 0000000000..f3547ccc5c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_256_64_kernel.sm80.cu
@@ -0,0 +1,58 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include "fmha.h"
+#include "fmha_fprop_kernel_1xN.h"
+
+using Kernel_traits = FMHA_kernel_traits< 256, 64, 16, 1, 4, 0x08u>;
+
+extern "C" __global__ void fmha_fprop_fp16_256_64_sm80_train_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::device_1xN<Kernel_traits, true>(params);
+}
+
+extern "C" __global__ void fmha_fprop_fp16_256_64_sm80_predict_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::device_1xN<Kernel_traits, false>(params);
+}
+
+void run_fmha_fp16_256_64_sm80(const Fused_multihead_attention_fprop_params &params, bool is_training, cudaStream_t stream) {
+
+    auto kernel = is_training ? &fmha_fprop_fp16_256_64_sm80_train_kernel : &fmha_fprop_fp16_256_64_sm80_predict_kernel;
+
+    constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);
+    constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;
+    constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;
+    constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;
+
+    constexpr int smem_size = smem_size_q + std::max(smem_size_v, smem_size_o + smem_size_softmax);
+
+    if( smem_size >= 48 * 1024 ) {
+        FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+
+    dim3 grid(params.h, params.b);
+    kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_384_64_kernel.sm80.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_384_64_kernel.sm80.cu
new file mode 100644
index 0000000000..f29851dfc4
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_384_64_kernel.sm80.cu
@@ -0,0 +1,57 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include "fmha.h"
+#include "fmha_fprop_kernel_1xN_reload_v.h"
+
+using Kernel_traits = FMHA_kernel_traits< 384, 64, 16, 1, 4, 0x08u>;
+
+extern "C" __global__ void fmha_fprop_fp16_384_64_sm80_train_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::device_1xN<Kernel_traits, true>(params);
+}
+
+extern "C" __global__ void fmha_fprop_fp16_384_64_sm80_predict_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::device_1xN<Kernel_traits, false>(params);
+}
+
+void run_fmha_fp16_384_64_sm80(const Fused_multihead_attention_fprop_params &params, bool is_training, cudaStream_t stream) {
+
+    auto kernel = is_training ? &fmha_fprop_fp16_384_64_sm80_train_kernel : &fmha_fprop_fp16_384_64_sm80_predict_kernel;
+
+    constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);
+    constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;
+    constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;
+
+    constexpr int smem_size = smem_size_v + smem_size_o + smem_size_softmax;
+
+    if( smem_size >= 48 * 1024 ) {
+        FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+
+    dim3 grid(params.h, params.b);
+    kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_512_64_kernel.sm80.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_512_64_kernel.sm80.cu
new file mode 100644
index 0000000000..4d5f259f4b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_fp16_512_64_kernel.sm80.cu
@@ -0,0 +1,98 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include "fmha.h"
+#include "fmha_fprop_kernel_1xN.h"
+#include "fmha_fprop_kernel_1xN_nl.h"
+
+using Kernel_traits = FMHA_kernel_traits< 512, 64, 16, 1, 8, 0x08u>;
+
+extern "C" __global__ void fmha_fprop_fp16_512_64_sm80_train_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::device_1xN<Kernel_traits, true>(params);
+}
+
+extern "C" __global__ void fmha_fprop_fp16_512_64_sm80_predict_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::device_1xN<Kernel_traits, false>(params);
+}
+
+template<int CHUNKS>
+__global__ void fmha_fprop_fp16_512_64_sm80_train_nl_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::device_1xN_nl<CHUNKS,Kernel_traits, true>(params);
+}
+
+template<int CHUNKS>
+__global__ void fmha_fprop_fp16_512_64_sm80_predict_nl_kernel(Fused_multihead_attention_fprop_params params) {
+    fmha::device_1xN_nl<CHUNKS, Kernel_traits, false>(params);
+}
+
+
+void run_fmha_fp16_512_64_sm80(const Fused_multihead_attention_fprop_params &params, bool is_training, cudaStream_t stream) {
+
+    auto kernel = is_training ? &fmha_fprop_fp16_512_64_sm80_train_kernel : &fmha_fprop_fp16_512_64_sm80_predict_kernel;
+
+    constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);
+    constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;
+    constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;
+    constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;
+
+    constexpr int smem_size = smem_size_q + std::max(smem_size_v, smem_size_o + smem_size_softmax);
+    if( smem_size >= 48 * 1024 ) {
+        FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+    dim3 grid(params.h, params.b);
+    kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);
+}
+
+void run_fmha_fp16_512_64_sm80_nl(const Fused_multihead_attention_fprop_params &params, const bool is_training, const int num_chunks, cudaStream_t stream) {
+
+    auto kernel = is_training ? &fmha_fprop_fp16_512_64_sm80_train_nl_kernel<2> : &fmha_fprop_fp16_512_64_sm80_predict_nl_kernel<2>;
+    if( num_chunks == 2 ) {
+        kernel = is_training ? &fmha_fprop_fp16_512_64_sm80_train_nl_kernel<2>
+                             : &fmha_fprop_fp16_512_64_sm80_predict_nl_kernel<2>;
+    } else if( num_chunks == 3 ) {
+        kernel = is_training ? &fmha_fprop_fp16_512_64_sm80_train_nl_kernel<3>
+                             : &fmha_fprop_fp16_512_64_sm80_predict_nl_kernel<3>;
+    } else if( num_chunks == 4 ) {
+        kernel = is_training ? &fmha_fprop_fp16_512_64_sm80_train_nl_kernel<4>
+                             : &fmha_fprop_fp16_512_64_sm80_predict_nl_kernel<4>;
+    } else {
+        assert(false && "Unsupported num_chunks");
+    }
+
+    constexpr int smem_size_softmax = Kernel_traits::Cta_tile_p::M * Kernel_traits::Cta_tile_p::WARPS_N * sizeof(float);
+    constexpr int smem_size_q = Kernel_traits::Smem_tile_q::BYTES_PER_TILE;
+    constexpr int smem_size_v = Kernel_traits::Smem_tile_v::BYTES_PER_TILE;
+    constexpr int smem_size_o = Kernel_traits::Smem_tile_o::BYTES_PER_TILE;
+
+    constexpr int smem_size = smem_size_q + std::max(smem_size_v, smem_size_o + smem_size_softmax);
+    if( smem_size >= 48 * 1024 ) {
+        FMHA_CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, smem_size));
+    }
+
+    dim3 grid(params.h, params.b, num_chunks);
+    kernel<<<grid, Kernel_traits::THREADS, smem_size, stream>>>(params);
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN.h
new file mode 100644
index 0000000000..ed8d92ac64
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN.h
@@ -0,0 +1,336 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include "fmha_kernel.h"
+#include <fmha/kernel_traits.h>
+#include <fmha/gemm.h>
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Kernel_traits, bool Is_training, typename Params> inline __device__ void device_1xN(const Params &params) {
+
+    // The description of the CTA tile for the 1st batched GEMM.
+    using Cta_tile_p = typename Kernel_traits::Cta_tile_p;
+    // The description of the CTA tile for the 2nd batched GEMM.
+    using Cta_tile_o = typename Kernel_traits::Cta_tile_o;
+
+    // The MMA tile for the 1st GEMM.
+    using Mma_tile_p = fmha::Hmma_tile<Cta_tile_p>;
+    // The MMA tile for the 2nd GEMM.
+    using Mma_tile_o = fmha::Hmma_tile<Cta_tile_o>;
+
+    // The global memory tile to load Q.
+    using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;
+    // The shared memory tile to swizzle Q.
+    using Smem_tile_q = typename Kernel_traits::Smem_tile_q;
+
+    // The global memory tile to load K.
+    using Gmem_tile_k = typename Kernel_traits::Gmem_tile_k;
+    // The shared memory tile to swizzle K.
+    using Smem_tile_k = typename Kernel_traits::Smem_tile_k;
+
+    // The global memory tile to load V.
+    using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle V.
+    using Smem_tile_v = typename Kernel_traits::Smem_tile_v;
+
+    // The global memory tile to store O.
+    using Gmem_tile_o = typename Kernel_traits::Gmem_tile_o;
+    // The shared memory tile to swizzle O.
+    using Smem_tile_o = typename Kernel_traits::Smem_tile_o;
+
+    using Gmem_tile_s = typename Kernel_traits::Gmem_tile_s;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+
+    // The block index for the batch.
+    const int bidb = blockIdx.y;
+    // The block index for the head.
+    const int bidh = blockIdx.x;
+    // The thread index.
+    const int tidx = threadIdx.x;
+
+    const BlockInfoPadded<Kernel_traits::THREADS> binfo(params, bidb, bidh, tidx);
+    if( binfo.stop_early() )
+        return;
+
+    auto seeds = at::cuda::philox::unpack(params.philox_args);
+
+    Philox ph(std::get<0>(seeds), binfo.tidx_global, std::get<1>(seeds));
+
+    Mask<Cta_tile_p> mask(params, binfo, tidx);
+
+    // Allocate the global memory tile loader for Q.
+    Gmem_tile_q gmem_q(params, 0, binfo, tidx);
+    // Allocate the shared memory tile loader for Q.
+    Smem_tile_q smem_q(&smem_[0], tidx);
+
+    // Allocate the global memory tile loader for K.
+    Gmem_tile_k gmem_k(params, 1, binfo, tidx);
+    // Allocate the shared memory tile loader for K.
+    Smem_tile_k smem_k(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);
+
+    // Allocate the global memory tile loader for V.
+    Gmem_tile_v gmem_v(params, 2, binfo, tidx);
+    // The base pointer of smem_v;
+    char *smem_v_ = nullptr;
+    if( Kernel_traits::SHARE_SMEM_FOR_K_AND_V ) {
+        smem_v_ = &smem_[Smem_tile_q::BYTES_PER_TILE];
+    } else {
+        smem_v_ = &smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE];
+    }
+    // Allocate the shared memory tile loader for V. We use the same as K so be careful!!!
+    Smem_tile_v smem_v(smem_v_, tidx);
+
+    // Allocate the global memory tile loader for O.
+    Gmem_tile_o gmem_o(params, binfo, tidx);
+    // Allocate the shared memory tile loader for O. We use the same as K so be careful!!!
+    Smem_tile_o smem_o(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);
+
+    // Trigger the loads for Q.
+    gmem_q.load(smem_q);
+    // Trigger the loads for K.
+    gmem_k.load(smem_k);
+    // Trigger the loads for K.
+    gmem_v.load(smem_v);
+
+    // Commit the data for Q and K to shared memory.
+    gmem_q.commit(smem_q);
+    gmem_k.commit(smem_k);
+
+    // Commit the data for V to shared memory.
+    if( !Kernel_traits::SHARE_SMEM_FOR_K_AND_V ) {
+        gmem_v.commit(smem_v);
+    }
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Load the fragments for Q.
+    typename Smem_tile_q::Fragment frag_q[2][Mma_tile_p::MMAS_M];
+    smem_q.load(frag_q[0], 0);
+
+    // Load the fragments for K. We keep the data in registers during the entire kernel.
+    typename Smem_tile_k::Fragment frag_k[Mma_tile_p::MMAS_K][Mma_tile_p::MMAS_N];
+    #pragma unroll
+    for( int ki = 0; ki < Mma_tile_p::MMAS_K; ++ki ) {
+        smem_k.load(frag_k[ki], ki);
+    }
+
+    // Commit the data for V to shared memory if it has not been done already.
+    if( Kernel_traits::SHARE_SMEM_FOR_K_AND_V ) {
+        // Make sure we are done loading the fragments for K.
+        __syncthreads();
+
+        // Commit the data to shared memory for V.
+        gmem_v.commit(smem_v);
+
+        // Make sure the data is in shared memory.
+        __syncthreads();
+    }
+
+    // Load the fragments for V. We keep the data in registers during the entire kernel.
+    typename Smem_tile_v::Fragment frag_v[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_N];
+    #pragma unroll
+    for( int ki = 0; ki < Mma_tile_o::MMAS_K; ++ki ) {
+        smem_v.load(frag_v[ki], ki);
+    }
+
+    enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };
+
+    Gmem_tile_s gmem_s(params.s_ptr, params, tidx);
+
+    // Create the object to do the softmax.
+    using Softmax = fmha::Softmax< Cta_tile_p, Kernel_traits>;
+    Softmax softmax(params, &smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE], bidb, tidx);
+
+    enum { THREADS_PER_ROW = 32 };
+    enum { STEPS = Cta_tile_p::N / Cta_tile_p::M };
+
+    // Load over the entire sequence length.
+    for( int l = 0; l < STEPS; l++ ) {
+        const int loop = l * Cta_tile_p::M;
+        if( loop >= binfo.actual_seqlen )
+            break;
+
+        // Declare the accumulators for the 1st gemm.
+        fmha::Fragment_accumulator acc_p[Mma_tile_p::MMAS_M][Mma_tile_p::MMAS_N];
+        fmha::Clear_accumulator<typename fmha::Accumulator_type, Cta_tile_p::WARPS_K>::apply(acc_p);
+
+    // Do this part of P^T = (Q * K^T)^T.
+    #pragma unroll
+        for( int ki = 1; ki < Mma_tile_p::MMAS_K; ++ki ) {
+
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_q.load(frag_q[ki & 1], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1)]);
+        }
+
+        // Do the final stage of math.
+        {
+            int ki = Mma_tile_p::MMAS_K;
+            fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1)]);
+        }
+
+        // Load the mask for that iteration.
+        mask.load(l);
+
+        // Convert from the accumulator type to FP32 for Softmax.
+        softmax.unpack(acc_p);
+
+        // Apply the mask.
+        softmax.apply_mask(mask);
+
+        if( Kernel_traits::SHARE_SMEM_FOR_K_AND_V && l == 0 ) {
+            // if we share K and V, it could be that V was not fully read yet but we write into smem for reduction
+            __syncthreads();
+        }
+        // Compute the max.
+        float p_max[Mma_tile_p::MMAS_M * 2];
+        softmax.template reduce<fmha::Max_>(p_max);
+
+        // Make sure we are done reading shared memory.
+        __syncthreads();
+
+        // Compute the exponential value.
+        softmax.apply_exp(p_max);
+
+        // Compute the sum.
+        float p_sum[Mma_tile_p::MMAS_M * 2];
+        softmax.template reduce<fmha::Sum_>(p_sum);
+
+        // Finalize softmax on the accumulators of P^T.
+        softmax.scale(p_sum);
+
+        if( Is_training ) {
+            auto encode_dropout = [](bool keep, float val) { return keep ? val : -val; };
+            #pragma unroll
+            for( int mi = 0; mi < Mma_tile_p::MMAS_M; mi++ ) {
+                #pragma unroll
+                for( int ii = 0; ii < 2; ii++ ) {
+                    #pragma unroll
+                    for( int ni = 0; ni < Mma_tile_p::MMAS_N; ni++ ) {
+                        float4 tmp = uniform4(ph());
+                        // We encode the dropout pattern in the sign bit of the non-negative softmax to distinguish from
+                        // pre-existing zeros
+                        softmax.elt_[2 * mi + ii][4 * ni + 0] =
+                            encode_dropout(tmp.x <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 0]);
+                        softmax.elt_[2 * mi + ii][4 * ni + 1] =
+                            encode_dropout(tmp.y <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 1]);
+                        softmax.elt_[2 * mi + ii][4 * ni + 2] =
+                            encode_dropout(tmp.z <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 2]);
+                        softmax.elt_[2 * mi + ii][4 * ni + 3] =
+                            encode_dropout(tmp.w <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 3]);
+                    }
+                }
+            }
+            gmem_s.store(softmax.elt_, mask);
+            gmem_s.move();
+        }
+
+        // Trigger the load for the next Q values.
+        if(l < STEPS - 1) {
+            smem_q.move_to_next_write_buffer();
+            gmem_q.move();
+            gmem_q.load(smem_q);
+        }
+
+        using Frag_p = fmha::Fragment_a< fmha::Row>;
+        Frag_p frag_p[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_M];
+        softmax.pack(frag_p);
+        #pragma unroll
+        for( int ki = 0; ki < Mma_tile_o::MMAS_K; ki++ ) {
+            #pragma unroll
+            for( int mi = 0; mi < Mma_tile_o::MMAS_M; mi++ ) {
+                #pragma unroll
+                for( int ii = 0; ii < Frag_p::NUM_REGS; ii++ ) {
+                    //"Apply" the dropout.
+                    frag_p[ki][mi].reg(ii) = fmha::hmul2(frag_p[ki][mi].reg(ii), params.scale_dropout);
+                    frag_p[ki][mi].reg(ii) = fmha::hrelu2(frag_p[ki][mi].reg(ii));
+                }
+            }
+        }
+
+        // Declare the accumulators for the 1st gemm.
+        fmha::Fragment_accumulator acc_o[Mma_tile_o::MMAS_M][Mma_tile_o::MMAS_N];
+        fmha::Clear_accumulator<typename fmha::Accumulator_type, Cta_tile_o::WARPS_K>::apply(acc_o);
+
+        // Do this part of O = P^T * V^T.
+        #pragma unroll
+        for( int ki = 0; ki < Mma_tile_o::MMAS_K; ++ki ) {
+            fmha::gemm(acc_o, frag_p[ki], frag_v[ki]);
+        }
+
+        // Loop over MMAS_M.
+        #pragma unroll
+        for( int ii = 0; ii < Gmem_tile_o::LOOPS; ++ii ) {
+
+            // Swizzle the elements and do the final reduction.
+            smem_o.store(acc_o, ii);
+
+            // Make sure the data is in shared memory.
+            __syncthreads();
+
+            // Load from shared memory.
+            uint4 out[Gmem_tile_o::STGS_PER_LOOP];
+            smem_o.load(out);
+
+            // Make sure the data was read from shared memory.
+            if( ii < Gmem_tile_o::LOOPS - 1 ) {
+                __syncthreads();
+            }
+
+            // Output the values.
+            gmem_o.store(out, ii);
+        }
+
+        // Move to the next part of the output.
+        gmem_o.move();
+
+        // Commit the values for Q into shared memory.
+        if(l < STEPS - 1) {
+            gmem_q.commit(smem_q);
+        }
+
+        // Make sure the data is in shared memory.
+        __syncthreads();
+
+        // Trigger the loads for the values of Q for the next iteration.
+        smem_q.load(frag_q[0], 0);
+
+    }  // Outer loop over the sequence length.
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN_nl.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN_nl.h
new file mode 100644
index 0000000000..82fffd75de
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN_nl.h
@@ -0,0 +1,343 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include "fmha.h"
+#include <fmha/kernel_traits.h>
+#include <fmha/gemm.h>
+
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int CHUNKS, typename Kernel_traits, bool Is_training, typename Params>
+inline __device__ void device_1xN_nl(const Params &params) {
+
+    // The description of the CTA tile for the 1st batched GEMM.
+    using Cta_tile_p = typename Kernel_traits::Cta_tile_p;
+    // The description of the CTA tile for the 2nd batched GEMM.
+    using Cta_tile_o = typename Kernel_traits::Cta_tile_o;
+
+    // The MMA tile for the 1st GEMM.
+    using Mma_tile_p = fmha::Hmma_tile<Cta_tile_p>;
+    // The MMA tile for the 2nd GEMM.
+    using Mma_tile_o = fmha::Hmma_tile<Cta_tile_o>;
+
+    // The global memory tile to load Q.
+    using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;
+    // The shared memory tile to swizzle Q.
+    using Smem_tile_q = typename Kernel_traits::Smem_tile_q;
+
+    // The global memory tile to load K.
+    using Gmem_tile_k = typename Kernel_traits::Gmem_tile_k;
+    // The shared memory tile to swizzle K.
+    using Smem_tile_k = typename Kernel_traits::Smem_tile_k;
+
+    // The global memory tile to load V.
+    using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle V.
+    using Smem_tile_v = typename Kernel_traits::Smem_tile_v;
+
+    // The global memory tile to store O.
+    using Gmem_tile_o = typename Kernel_traits::Gmem_tile_o;
+    // The shared memory tile to swizzle O.
+    using Smem_tile_o = typename Kernel_traits::Smem_tile_o;
+
+    // The global memory tile to store S/D.
+    using Gmem_tile_s = typename Kernel_traits::Gmem_tile_s;
+
+    using Noloop = Noloop_traits<CHUNKS, Cta_tile_p>;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+
+    const int bidc = blockIdx.z;
+    // The block index for the batch.
+    const int bidb = blockIdx.y;
+    // The block index for the head.
+    const int bidh = blockIdx.x;
+    // The thread index.
+    const int tidx = threadIdx.x;
+
+    Noloop nl_traits(bidc);
+
+    const BlockInfoPadded<Kernel_traits::THREADS> binfo(params, bidb, bidh, tidx);
+    if( binfo.stop_early() )
+        return;
+
+    auto seeds = at::cuda::philox::unpack(params.philox_args);
+
+    Philox ph(std::get<0>(seeds), binfo.tidx_global, std::get<1>(seeds));
+
+    fmha::Mask<Cta_tile_p> mask(params, binfo, tidx);
+
+    // Allocate the global memory tile loader for Q.
+    Gmem_tile_q gmem_q(params, 0, binfo, tidx);
+    // Allocate the shared memory tile loader for Q.
+    Smem_tile_q smem_q(&smem_[0], tidx);
+
+    // Allocate the global memory tile loader for K.
+    Gmem_tile_k gmem_k(params, 1, binfo, tidx);
+    // Allocate the shared memory tile loader for K.
+    Smem_tile_k smem_k(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);
+
+    // Allocate the global memory tile loader for V.
+    Gmem_tile_v gmem_v(params, 2, binfo, tidx);
+    // The base pointer of smem_v;
+    char *smem_v_ = nullptr;
+    if( Kernel_traits::SHARE_SMEM_FOR_K_AND_V ) {
+        smem_v_ = &smem_[Smem_tile_q::BYTES_PER_TILE];
+    } else {
+        smem_v_ = &smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE];
+    }
+    // Allocate the shared memory tile loader for V. We use the same as K so be careful!!!
+    Smem_tile_v smem_v(smem_v_, tidx);
+
+    // Allocate the global memory tile loader for O.
+    Gmem_tile_o gmem_o(params, binfo, tidx);
+    // Allocate the shared memory tile loader for O. We use the same as K so be careful!!!
+    Smem_tile_o smem_o(&smem_[Smem_tile_q::BYTES_PER_TILE], tidx);
+
+    Gmem_tile_s gmem_s(params.s_ptr, params, tidx);
+
+    nl_traits.move_all(gmem_q, gmem_o, gmem_s);
+
+
+    // Trigger the loads for Q.
+    gmem_q.load(smem_q);
+    // Trigger the loads for K.
+    gmem_k.load(smem_k);
+    // Trigger the loads for K.
+    gmem_v.load(smem_v);
+
+
+    // Commit the data for Q and K to shared memory.
+    gmem_q.commit(smem_q);
+    gmem_k.commit(smem_k);
+
+    // Commit the data for V to shared memory.
+    if( !Kernel_traits::SHARE_SMEM_FOR_K_AND_V ) {
+        gmem_v.commit(smem_v);
+    }
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Load the fragments for Q.
+    typename Smem_tile_q::Fragment frag_q[2][Mma_tile_p::MMAS_M];
+    smem_q.load(frag_q[0], 0);
+
+    // Load the fragments for K. We keep the data in registers during the entire kernel.
+    typename Smem_tile_k::Fragment frag_k[Mma_tile_p::MMAS_K][Mma_tile_p::MMAS_N];
+    #pragma unroll
+    for( int ki = 0; ki < Mma_tile_p::MMAS_K; ++ki ) {
+        smem_k.load(frag_k[ki], ki);
+    }
+
+    // Commit the data for V to shared memory if it has not been done already.
+    if( Kernel_traits::SHARE_SMEM_FOR_K_AND_V ) {
+        // Make sure we are done loading the fragments for K.
+        __syncthreads();
+
+        // Commit the data to shared memory for V.
+        gmem_v.commit(smem_v);
+
+        // Make sure the data is in shared memory.
+        __syncthreads();
+    }
+
+    // Load the fragments for V. We keep the data in registers during the entire kernel.
+    typename Smem_tile_v::Fragment frag_v[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_N];
+    #pragma unroll
+    for( int ki = 0; ki < Mma_tile_o::MMAS_K; ++ki ) {
+        smem_v.load(frag_v[ki], ki);
+    }
+
+    enum { BITS_PER_ELT_S = sizeof(fmha::A_type) * 8 };
+
+
+    // Create the object to do the softmax.
+    using Softmax = fmha::Softmax<Cta_tile_p, Kernel_traits>;
+    Softmax softmax(params, &smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE], bidb, tidx);
+
+    // The number of threads per row.
+    enum { THREADS_PER_ROW = 32 };
+
+    // Load over the entire sequence length.
+    for(int l = 0; l < nl_traits.num_steps_;l++) {
+
+        // Declare the accumulators for the 1st gemm.
+        fmha::Fragment_accumulator acc_p[Mma_tile_p::MMAS_M][Mma_tile_p::MMAS_N];
+        fmha::Clear_accumulator<typename fmha::Accumulator_type, Cta_tile_p::WARPS_K>::apply(acc_p);
+
+        // Do this part of P^T = (Q * K^T)^T.
+        #pragma unroll
+        for( int ki = 1; ki < Mma_tile_p::MMAS_K; ++ki ) {
+
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_q.load(frag_q[ki & 1], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1)]);
+        }
+
+        // Do the final stage of math.
+        {
+            int ki = Mma_tile_p::MMAS_K;
+            fmha::gemm(acc_p, frag_q[(ki - 1) & 1], frag_k[(ki - 1)]);
+        }
+
+        // Trigger the load for the next Q values.
+        if( l < nl_traits.num_steps_- 1) {
+            smem_q.move_to_next_write_buffer();
+            gmem_q.move();
+            gmem_q.load(smem_q);
+        }
+
+
+
+        // Load the mask for that iteration.
+        mask.load(nl_traits.loop_offset_ + l);
+
+        // Convert from the accumulator type to FP32 for Softmax.
+        softmax.unpack(acc_p);
+
+        // Apply the mask.
+        softmax.apply_mask(mask);
+
+        if( Kernel_traits::SHARE_SMEM_FOR_K_AND_V && l == 0 ) {
+            // if we share K and V, it could be that V was not fully read yet but we write into smem for reduction
+            __syncthreads();
+        }
+
+        // Compute the max.
+        float p_max[Mma_tile_p::MMAS_M * 2];
+        softmax.template reduce<fmha::Max_>(p_max);
+
+        // Make sure we are done reading shared memory.
+        __syncthreads();
+
+        // Compute the exponential value.
+        softmax.apply_exp(p_max);
+
+        // Compute the sum.
+        float p_sum[Mma_tile_p::MMAS_M * 2];
+        softmax.template reduce<fmha::Sum_>(p_sum);
+
+        // Finalize softmax on the accumulators of P^T.
+        softmax.scale(p_sum);
+        if( Is_training ) {
+            auto encode_dropout = [](bool keep, float val) { return keep ? val : -val; };
+            #pragma unroll
+            for( int mi = 0; mi < Mma_tile_p::MMAS_M; mi++ ) {
+                #pragma unroll
+                for( int ii = 0; ii < 2; ii++ ) {
+                    #pragma unroll
+                    for( int ni = 0; ni < Mma_tile_p::MMAS_N; ni++ ) {
+                        float4 tmp = uniform4(ph());
+                        // We encode the dropout pattern in the sign bit of the non-negative softmax to distinguish from pre-existing zeros
+                        softmax.elt_[2 * mi + ii][4 * ni + 0] =
+                            encode_dropout(tmp.x <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 0]);
+                        softmax.elt_[2 * mi + ii][4 * ni + 1] =
+                            encode_dropout(tmp.y <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 1]);
+                        softmax.elt_[2 * mi + ii][4 * ni + 2] =
+                            encode_dropout(tmp.z <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 2]);
+                        softmax.elt_[2 * mi + ii][4 * ni + 3] =
+                            encode_dropout(tmp.w <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 3]);
+                    }
+                }
+            }
+            gmem_s.store(softmax.elt_, mask);
+            gmem_s.move();
+        }
+
+        using Frag_p = fmha::Fragment_a<fmha::Row>;
+        Frag_p frag_p[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_M];
+        softmax.pack(frag_p);
+        #pragma unroll
+        for( int ki = 0; ki < Mma_tile_o::MMAS_K; ki++ ) {
+            #pragma unroll
+            for( int mi = 0; mi < Mma_tile_o::MMAS_M; mi++ ) {
+                #pragma unroll
+                for( int ii = 0; ii < Frag_p::NUM_REGS; ii++ ) {
+                    //"Apply" the dropout.
+                    frag_p[ki][mi].reg(ii) = fmha::hmul2(frag_p[ki][mi].reg(ii), params.scale_dropout);
+                    frag_p[ki][mi].reg(ii) = fmha::hrelu2(frag_p[ki][mi].reg(ii));
+                }
+            }
+        }
+
+        // Declare the accumulators for the 1st gemm.
+        fmha::Fragment_accumulator acc_o[Mma_tile_o::MMAS_M][Mma_tile_o::MMAS_N];
+        fmha::Clear_accumulator<typename fmha::Accumulator_type, Cta_tile_o::WARPS_K>::apply(acc_o);
+
+        // Do this part of O = P^T * V^T.
+        #pragma unroll
+        for( int ki = 0; ki < Mma_tile_o::MMAS_K; ++ki ) {
+            fmha::gemm(acc_o, frag_p[ki], frag_v[ki]);
+        }
+
+        // Loop over MMAS_M.
+        #pragma unroll
+        for( int ii = 0; ii < Gmem_tile_o::LOOPS; ++ii ) {
+
+            // Swizzle the elements and do the final reduction.
+            smem_o.store(acc_o, ii);
+
+            // Make sure the data is in shared memory.
+            __syncthreads();
+
+            // Load from shared memory.
+            uint4 out[Gmem_tile_o::STGS_PER_LOOP];
+            smem_o.load(out);
+
+            // Make sure the data was read from shared memory.
+            if( ii < Gmem_tile_o::LOOPS - 1 ) {
+                __syncthreads();
+            }
+
+            // Output the values.
+            gmem_o.store(out, ii);
+        }
+
+        // Move to the next part of the output.
+        gmem_o.move();
+
+        // Commit the values for Q into shared memory.
+        if( l < nl_traits.num_steps_- 1) {
+            gmem_q.commit(smem_q);
+            __syncthreads();
+            smem_q.load(frag_q[0], 0);
+        }
+
+    }  // Outer loop over the sequence length.
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+} // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN_reload_v.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN_reload_v.h
new file mode 100644
index 0000000000..50fd43f6be
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_fprop_kernel_1xN_reload_v.h
@@ -0,0 +1,322 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include "fmha_kernel.h"
+#include <fmha/kernel_traits.h>
+#include <fmha/gemm.h>
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Kernel_traits, bool Is_training, typename Params> inline __device__ void device_1xN(const Params &params) {
+
+    // The description of the CTA tile for the 1st batched GEMM.
+    using Cta_tile_p = typename Kernel_traits::Cta_tile_p;
+    // The description of the CTA tile for the 2nd batched GEMM.
+    using Cta_tile_o = typename Kernel_traits::Cta_tile_o;
+
+    // The MMA tile for the 1st GEMM.
+    using Mma_tile_p = fmha::Hmma_tile<Cta_tile_p>;
+    // The MMA tile for the 2nd GEMM.
+    using Mma_tile_o = fmha::Hmma_tile<Cta_tile_o>;
+
+    // The global memory tile to load Q.
+    using Gmem_tile_q = typename Kernel_traits::Gmem_tile_q;
+    // The shared memory tile to swizzle Q.
+    using Smem_tile_q = typename Kernel_traits::Smem_tile_q;
+
+    // The global memory tile to load K.
+    using Gmem_tile_k = typename Kernel_traits::Gmem_tile_k;
+    // The shared memory tile to swizzle K.
+    using Smem_tile_k = typename Kernel_traits::Smem_tile_k;
+
+    // The global memory tile to load V.
+    using Gmem_tile_v = typename Kernel_traits::Gmem_tile_v;
+    // The shared memory tile to swizzle V.
+    using Smem_tile_v = typename Kernel_traits::Smem_tile_v;
+
+    // The global memory tile to store O.
+    using Gmem_tile_o = typename Kernel_traits::Gmem_tile_o;
+    // The shared memory tile to swizzle O.
+    using Smem_tile_o = typename Kernel_traits::Smem_tile_o;
+
+    using Gmem_tile_s = typename Kernel_traits::Gmem_tile_s;
+
+    // Shared memory.
+    extern __shared__ char smem_[];
+
+    // The block index for the batch.
+    const int bidb = blockIdx.y;
+    // The block index for the head.
+    const int bidh = blockIdx.x;
+    // The thread index.
+    const int tidx = threadIdx.x;
+
+    const BlockInfoPadded<Kernel_traits::THREADS> binfo(params, bidb, bidh, tidx);
+    if( binfo.stop_early() )
+        return;
+
+    Mask<Cta_tile_p> mask(params, binfo, tidx);
+
+    auto seeds = at::cuda::philox::unpack(params.philox_args);
+    Philox ph(std::get<0>(seeds), binfo.tidx_global, std::get<1>(seeds));
+
+    static_assert(2 * Mma_tile_p::MMAS_M * 4 * Mma_tile_p::MMAS_N <= 64);
+
+    // Allocate the global memory tile loader for K.
+    Gmem_tile_k gmem_k(params, 1, binfo, tidx);
+    // Allocate the shared memory tile loader for K.
+    Smem_tile_k smem_k(&smem_[0], tidx);
+
+    // Allocate the global memory tile loader for V.
+    Gmem_tile_v gmem_v(params, 2, binfo, tidx);
+    // The base pointer of smem_v;
+    char *smem_v_ = nullptr;
+    if( Kernel_traits::SHARE_SMEM_FOR_K_AND_V ) {
+        smem_v_ = &smem_[0];
+    } else {
+        smem_v_ = &smem_[Smem_tile_q::BYTES_PER_TILE + Smem_tile_k::BYTES_PER_TILE];
+    }
+    static_assert(Kernel_traits::SHARE_SMEM_FOR_K_AND_V);
+    static_assert(Smem_tile_k::BYTES_PER_TILE == Smem_tile_v::BYTES_PER_TILE);
+    // Allocate the shared memory tile loader for V. We use the same as K so be careful!!!
+    Smem_tile_v smem_v(smem_v_, tidx);
+
+    // Allocate the global memory tile loader for Q.
+    Gmem_tile_q gmem_q(params, 0, binfo, tidx);
+    // Allocate the shared memory tile loader for Q.
+    Smem_tile_q smem_q(&smem_[Smem_tile_v::BYTES_PER_TILE], tidx);
+
+    // Allocate the global memory tile loader for O.
+    Gmem_tile_o gmem_o(params, binfo, tidx);
+    // Allocate the shared memory tile loader for O. We use the same as K so be careful!!!
+    Smem_tile_o smem_o(&smem_[Smem_tile_v::BYTES_PER_TILE], tidx);
+
+    // Trigger the loads for Q.
+    gmem_q.load(smem_q);
+    // Trigger the loads for K.
+    gmem_k.load(smem_k);
+    // Trigger the loads for K.
+    gmem_v.load(smem_v);
+
+
+    // Commit the data for Q and K to shared memory.
+    gmem_q.commit(smem_q);
+    gmem_k.commit(smem_k);
+
+    // Commit the data for V to shared memory.
+    if( !Kernel_traits::SHARE_SMEM_FOR_K_AND_V ) {
+        gmem_v.commit(smem_v);
+    }
+
+    // Make sure the data is in shared memory.
+    __syncthreads();
+
+    // Load the fragments for Q.
+    typename Smem_tile_q::Fragment frag_q[1][Mma_tile_p::MMAS_M];
+
+    // Load the fragments for K. We keep the data in registers during the entire kernel.
+    typename Smem_tile_k::Fragment frag_k[Mma_tile_p::MMAS_K][Mma_tile_p::MMAS_N];
+    #pragma unroll
+    for( int ki = 0; ki < Mma_tile_p::MMAS_K; ++ki ) {
+        smem_k.load(frag_k[ki], ki);
+    }
+
+    // Commit the data for V to shared memory if it has not been done already.
+    if( Kernel_traits::SHARE_SMEM_FOR_K_AND_V ) {
+        // Make sure we are done loading the fragments for K.
+        __syncthreads();
+
+        // Commit the data to shared memory for V.
+        gmem_v.commit(smem_v);
+
+    }
+
+    enum { BITS_PER_ELT_S = sizeof(typename fmha::A_type) * 8 };
+
+    Gmem_tile_s gmem_s(params.s_ptr, params, tidx);
+
+    // Create the object to do the softmax.
+    using Softmax = fmha::Softmax< Cta_tile_p, Kernel_traits>;
+    Softmax softmax(params, &smem_[Smem_tile_v::BYTES_PER_TILE + Smem_tile_o::BYTES_PER_TILE], bidb, tidx);
+
+    constexpr int SMEM_BYTES_SOFTMAX = Softmax::ELEMENTS * sizeof(float);
+    static_assert(SMEM_BYTES_SOFTMAX == Cta_tile_p::M * Cta_tile_p::WARPS_N * sizeof(float));
+
+    enum { THREADS_PER_ROW = 32 };
+
+    const float pinv = 1.f / params.p_dropout;
+
+    // Load over the entire sequence length.
+    for( int loop = 0, outer = 0; loop < Cta_tile_p::N; loop += Cta_tile_p::M, outer++ ) {
+        if( loop >= binfo.actual_seqlen )
+            break;
+
+        // Declare the accumulators for the 1st gemm.
+        fmha::Fragment_accumulator acc_p[Mma_tile_p::MMAS_M][Mma_tile_p::MMAS_N];
+        fmha::Clear_accumulator<typename fmha::Accumulator_type, Cta_tile_p::WARPS_K>::apply(acc_p);
+        #pragma unroll
+        for( int ki = 0; ki < Mma_tile_p::MMAS_K; ++ki ) {
+
+            // Trigger the load from shared memory for the next series of Q values.
+            smem_q.load(frag_q[0], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_p, frag_q[0], frag_k[ki]);
+        }
+
+        // Load the mask for that iteration.
+        mask.load(outer);
+
+        // Convert from the accumulator typ e to FP32 for Softmax.
+        softmax.unpack(acc_p);
+
+        // Apply the mask.
+        softmax.apply_mask(mask);
+
+        static_assert(2 * Mma_tile_p::MMAS_M * 4 * Mma_tile_p::MMAS_N <= 64);
+
+        // Compute the max.
+        float p_max[Mma_tile_p::MMAS_M * 2];
+        softmax.template reduce<fmha::Max_>(p_max);
+
+        // Make sure we are done reading shared memory.
+        __syncthreads();
+        // Compute the exponential value.
+        softmax.apply_exp(p_max);
+        // Compute the sum.
+        float p_sum[Mma_tile_p::MMAS_M * 2];
+        softmax.template reduce<fmha::Sum_>(p_sum);
+
+        // Finalize softmax on the accumulators of P^T.
+        softmax.scale(p_sum);
+
+        __syncthreads();
+        if( Is_training ) {
+            auto encode_dropout = [](bool keep, float val) { return keep ? val : -val; };
+            #pragma unroll
+            for( int mi = 0; mi < Mma_tile_p::MMAS_M; mi++ ) {
+                #pragma unroll
+                for( int ii = 0; ii < 2; ii++ ) {
+                    #pragma unroll
+                    for( int ni = 0; ni < Mma_tile_p::MMAS_N; ni++ ) {
+                        float4 tmp = uniform4(ph());
+                        // We encode the dropout pattern in the sign bit of the non-negative softmax to distinguish from
+                        // pre-existing zeros
+                        softmax.elt_[2 * mi + ii][4 * ni + 0] =
+                            encode_dropout(tmp.x <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 0]);
+                        softmax.elt_[2 * mi + ii][4 * ni + 1] =
+                            encode_dropout(tmp.y <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 1]);
+                        softmax.elt_[2 * mi + ii][4 * ni + 2] =
+                            encode_dropout(tmp.z <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 2]);
+                        softmax.elt_[2 * mi + ii][4 * ni + 3] =
+                            encode_dropout(tmp.w <= params.p_dropout, softmax.elt_[2 * mi + ii][4 * ni + 3]);
+                    }
+                }
+            }
+
+            gmem_s.store(softmax.elt_, mask);
+            gmem_s.move();
+        }
+
+        // Trigger the load for the next Q values.
+        if( loop + Cta_tile_p::M < Cta_tile_p::N ) {
+            smem_q.move_to_next_write_buffer();
+            gmem_q.move();
+            gmem_q.load(smem_q);
+        }
+        typename Smem_tile_v::Fragment frag_v[1][Mma_tile_o::MMAS_N];
+
+        using Frag_p = fmha::Fragment_a< fmha::Row>;
+        Frag_p frag_p[Mma_tile_o::MMAS_K][Mma_tile_o::MMAS_M];
+        softmax.pack(frag_p);
+        #pragma unroll
+        for( int ki = 0; ki < Mma_tile_o::MMAS_K; ki++ ) {
+            #pragma unroll
+            for( int mi = 0; mi < Mma_tile_o::MMAS_M; mi++ ) {
+                #pragma unroll
+                for( int ii = 0; ii < Frag_p::NUM_REGS; ii++ ) {
+                    //"Apply" the dropout.
+                    frag_p[ki][mi].reg(ii) = fmha::hmul2(frag_p[ki][mi].reg(ii), params.scale_dropout);
+                    frag_p[ki][mi].reg(ii) = fmha::hrelu2(frag_p[ki][mi].reg(ii));
+                }
+            }
+        }
+
+        // Declare the accumulators for the 1st gemm.
+        fmha::Fragment_accumulator acc_o[Mma_tile_o::MMAS_M][Mma_tile_o::MMAS_N];
+        fmha::Clear_accumulator<typename fmha::Accumulator_type, Cta_tile_o::WARPS_K>::apply(acc_o);
+
+        #pragma unroll
+        for( int ki = 0; ki < Mma_tile_o::MMAS_K; ++ki ) {
+            // Trigger the load from shared memory for the next series of V values.
+            smem_v.load(frag_v[0], ki);
+            // Do the math for the values already in registers.
+            fmha::gemm(acc_o, frag_p[ki], frag_v[0]);
+        }
+
+        // Loop over MMAS_M.
+        #pragma unroll
+        for( int ii = 0; ii < Gmem_tile_o::LOOPS; ++ii ) {
+
+            // Swizzle the elements and do the final reduction.
+            smem_o.store(acc_o, ii);
+
+            // Make sure the data is in shared memory.
+            __syncthreads();
+
+            // Load from shared memory.
+            uint4 out[Gmem_tile_o::STGS_PER_LOOP];
+            smem_o.load(out);
+
+            // Always sync after last iter: shared smem_q and smem_o!
+            __syncthreads();
+
+            // Output the values.
+            gmem_o.store(out, ii);
+        }
+        // same smem as o
+
+        // Move to the next part of the output.
+        gmem_o.move();
+
+        // Commit the values for Q into shared memory.
+        if( loop + Cta_tile_p::M < Cta_tile_p::N ) {
+            gmem_q.commit(smem_q);
+        }
+
+        // Make sure the data is in shared memory.
+        __syncthreads();
+
+    }  // Outer loop over the sequence length.
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_kernel.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_kernel.h
new file mode 100644
index 0000000000..465f783e78
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_kernel.h
@@ -0,0 +1,169 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include <multihead_attn/philox.h>
+
+#include <fmha.h>
+#include <fmha/utils.h>
+#include <fmha/smem_tile.h>
+#include <fmha/gmem_tile.h>
+#include <fmha/mask.h>
+#include <fmha/softmax.h>
+
+namespace fmha {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int THREADS_PER_CTA>
+struct BlockInfoPadded {
+
+    template<typename Params>
+    __device__ BlockInfoPadded(const Params &params,
+                               const int bidb,
+                               const int bidh,
+                               const int tidx)
+        : bidb(bidb), bidh(bidh), h(params.h) {
+
+        // The block index.
+        sum_s = params.cu_seqlens[bidb];
+        actual_seqlen = params.cu_seqlens[bidb + 1] - sum_s;
+        bidx = sum_s * params.h + bidh;
+
+        tidx_global = (bidb * params.h + bidh) * THREADS_PER_CTA + tidx;
+    }
+
+    __device__ bool stop_early() const {
+        return actual_seqlen == 0;
+    }
+
+    int actual_seqlen;
+    int bidx;
+    int sum_s;
+    int bidh;
+    int bidb;
+    int tidx_global;
+    int h;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int CHUNKS, typename Cta_tile> 
+struct Noloop_traits{
+    // Interpretation of Cta_tile dims, i.e. Cta_tile_p:
+    enum{ STEP = Cta_tile::M };
+    enum{ SEQLEN = Cta_tile::N };
+
+    // The size of the subsequence this CTA is processing
+    enum { SUBSEQ = SEQLEN / CHUNKS };
+    static_assert(SUBSEQ * CHUNKS == SEQLEN);
+
+    // The number of steps to process the subsequence
+    enum { NUM_STEPS = SUBSEQ / STEP };
+    static_assert(NUM_STEPS  * Cta_tile::M == SUBSEQ);
+
+    inline __device__ Noloop_traits(const int bidc) 
+        : loop_offset_(NUM_STEPS * bidc)
+        , bidc_(bidc) {
+    }
+
+    template<typename ... Tiles> 
+    inline __device__ void move_all(Tiles & ... tiles) const {
+        using expand_type = int[];
+        for( int s = 0; s < loop_offset_; s++ ) {
+            expand_type{ (tiles.move(), 0)... };
+        }
+    }
+
+    inline __device__ int get_idx_dk() const {
+        //return bidc_;
+        return bidc_ * 2 + 0;
+    }
+
+    inline __device__ int get_idx_dv() const {
+        //return CHUNKS + bidc_;
+        return bidc_ * 2 + 1;
+    }
+
+    inline __device__ int offset_loop_count(const int l) {
+        // convert loop counter to position in the outer sequence
+        return (loop_offset_ + l) * STEP;
+    }
+
+    const int loop_offset_;
+    const uint32_t bidc_;
+    const int num_steps_ = NUM_STEPS;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Cta_tile> 
+struct Noloop_traits<3, Cta_tile>{
+    // Interpretation of Cta_tile dims, i.e. Cta_tile_p:
+    enum{ STEP = Cta_tile::M };
+    enum{ SEQLEN = Cta_tile::N };
+
+    static_assert(STEP == 16 && SEQLEN == 512);
+
+    inline __device__ Noloop_traits(const int bidc)
+        : bidc_(bidc)
+        , num_steps_(bidc < 2 ? 11 : 10) 
+        , loop_offset_(bidc * 11) {
+    }
+
+    template<typename ... Tiles> 
+    inline __device__ void move_all(Tiles & ... tiles) const {
+        using expand_type = int[];
+        for( int s = 0; s < loop_offset_; s++ ) {
+            expand_type{ (tiles.move(), 0)... };
+        }
+    }
+
+    inline __device__ int get_idx_dk() const {
+        //return bidc_;
+        return bidc_ * 2 + 0;
+    }
+
+    inline __device__ int get_idx_dv() const {
+        //return CHUNKS + bidc_;
+        return bidc_ * 2 + 1;
+    }
+
+    inline __device__ int offset_loop_count(const int l) {
+        // convert loop counter to position in the outer sequence
+        return (loop_offset_ + l) * STEP;
+    }
+
+    const int loop_offset_;
+    const uint32_t bidc_;
+    const int  num_steps_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace fmha
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_noloop_reduce.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_noloop_reduce.cu
new file mode 100644
index 0000000000..8e4b9efc32
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_noloop_reduce.cu
@@ -0,0 +1,177 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#include "fmha.h"
+
+inline __device__ float4 ldg128(const void *ptr) {
+    return *static_cast<const float4 *>(ptr);
+}
+
+inline __device__ void stg128(void *ptr, const float4 &data) {
+    *static_cast<float4 *>(ptr) = data;
+}
+
+template<typename T, int THREADS, int HIDDEN_SIZE, int CHUNKS>
+__global__ __launch_bounds__(THREADS) void fmha_noloop_reduce_kernel(void *__restrict__ out,
+                                                                     const void *__restrict__ in,
+                                                                     const int *__restrict__ cu_seqlens,
+                                                                     const int batch_size) {
+
+    enum { BYTES_PER_LDG = 16 };
+    enum { NUM_ELTS = BYTES_PER_LDG / sizeof(T) };
+
+    // One CTA hidden vector for K and V
+    enum { BYTES_PER_ROW = HIDDEN_SIZE * sizeof(T) * 2 };
+    // The stride in bytes in dQKV
+    enum { OUT_STRIDE_BYTES = 3 * HIDDEN_SIZE * sizeof(T) };
+    // The offset in bytes in dQKV to the dKV part for non-interleaved heads
+    enum { OUT_OFFSET_KV_BYTES = HIDDEN_SIZE * sizeof(T) };
+
+    static_assert(BYTES_PER_ROW == HIDDEN_SIZE * 2 * sizeof(T)); 
+
+    // Size in bytes of the input tile
+    enum { BYTES_PER_TILE = CHUNKS * BYTES_PER_ROW };
+
+    enum { BYTES_PER_CTA = THREADS * BYTES_PER_LDG };
+
+    enum { LDGS = BYTES_PER_ROW / BYTES_PER_CTA };
+    static_assert(BYTES_PER_CTA * LDGS == BYTES_PER_ROW);
+
+    union Vec_t {
+        float4 raw;
+        T elt[NUM_ELTS];
+    };
+
+    // ZERO-OUT invalid positions in dQKV
+    const int total = cu_seqlens[batch_size];
+    if(blockIdx.x >= total){
+        enum { BYTES_PER_QKV_ROW = 3 * HIDDEN_SIZE * sizeof(T) };
+        enum { STGS = BYTES_PER_QKV_ROW / BYTES_PER_LDG };
+
+        const float4 zeros = make_float4(0.f, 0.f, 0.f, 0.f);
+
+        char *base_ptr = static_cast<char *>(out) + blockIdx.x * OUT_STRIDE_BYTES;
+
+        for(int tidx = threadIdx.x; tidx < STGS; tidx += THREADS){
+            stg128(base_ptr + tidx * BYTES_PER_LDG, zeros);
+        }
+
+        return;
+    }
+
+    // SETUP
+    const int offset_in = blockIdx.x * BYTES_PER_TILE + threadIdx.x * BYTES_PER_LDG;
+    const char *ptr_in = static_cast<const char *>(in) + offset_in;
+
+    const int offset_out = blockIdx.x * OUT_STRIDE_BYTES + threadIdx.x * BYTES_PER_LDG;
+    char *ptr_out = static_cast<char *>(out) + OUT_OFFSET_KV_BYTES + offset_out;
+
+    // LOAD
+
+    Vec_t local_in[CHUNKS][LDGS];
+
+    #pragma unroll
+    for( int c = 0; c < CHUNKS; c++ ) {
+        #pragma unroll
+        for( int l = 0; l < LDGS; l++ ) {
+            int offset = c * BYTES_PER_ROW + l * BYTES_PER_CTA;
+            local_in[c][l].raw = ldg128(ptr_in + offset);
+        }
+    }
+
+    // UNPACK
+    float acc[LDGS][NUM_ELTS];
+
+    #pragma unroll
+    for( int l = 0; l < LDGS; l++ ) {
+        #pragma unroll
+        for( int e = 0; e < NUM_ELTS; e++ ) {
+            acc[l][e] = float(local_in[0][l].elt[e]);
+        }
+    }
+
+    // COMPUTE
+    #pragma unroll
+    for( int c = 1; c < CHUNKS; c++ ) {
+        #pragma unroll
+        for( int l = 0; l < LDGS; l++ ) {
+            #pragma unroll
+            for( int e = 0; e < NUM_ELTS; e++ ) {
+                acc[l][e] += float(local_in[c][l].elt[e]);
+            }
+        }
+    }
+
+    // PACK
+    Vec_t local_out[LDGS];
+
+    #pragma unroll
+    for( int l = 0; l < LDGS; l++ ) {
+        #pragma unroll
+        for( int e = 0; e < NUM_ELTS; e++ ) {
+            local_out[l].elt[e] = T(acc[l][e]);
+        }
+    }
+
+    // STORE
+    #pragma unroll
+    for( int l = 0; l < LDGS; l++ ) {
+        const int offset = l * BYTES_PER_CTA;
+        stg128(ptr_out + offset, local_out[l].raw);
+    }
+}
+
+void fmha_run_noloop_reduce(void *out,
+                            const void *in,
+                            const int *cu_seqlens,
+                            const int hidden_size,
+                            const int batch_size,
+                            const int total,
+                            const int num_chunks,
+                            cudaStream_t stream) {
+
+    const int blocks = total;
+
+    if(hidden_size == 1024){
+
+        constexpr int HIDDEN_SIZE = 1024;
+        constexpr int THREADS = 256;
+
+        if( num_chunks == 2 ) {
+            fmha_noloop_reduce_kernel<half, THREADS, HIDDEN_SIZE, 2><<<blocks, THREADS, 0, stream>>>(out, in, cu_seqlens, batch_size);
+        } else if( num_chunks == 3 ) {
+            fmha_noloop_reduce_kernel<half, THREADS, HIDDEN_SIZE, 3><<<blocks, THREADS, 0, stream>>>(out, in, cu_seqlens, batch_size);
+        } else {
+            assert(false && "Unsupported num_chunks");
+        }
+
+    }else{
+        assert(false && "Unsupported hidden_size");
+    }
+
+    FMHA_CHECK_CUDA(cudaPeekAtLastError());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_utils.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_utils.h
new file mode 100644
index 0000000000..de07cc78e3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/fmha/src/fmha_utils.h
@@ -0,0 +1,92 @@
+/******************************************************************************
+ * Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+ * 
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *     * Redistributions of source code must retain the above copyright
+ *       notice, this list of conditions and the following disclaimer.
+ *     * Redistributions in binary form must reproduce the above copyright
+ *       notice, this list of conditions and the following disclaimer in the
+ *       documentation and/or other materials provided with the distribution.
+ *     * Neither the name of the NVIDIA CORPORATION nor the
+ *       names of its contributors may be used to endorse or promote products
+ *       derived from this software without specific prior written permission.
+ * 
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+ * ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+ * WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+ * DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+ * DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+ * (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+ * LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+ * ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+ * (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+ * SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+ *
+ ******************************************************************************/
+
+#pragma once
+
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <cuda_runtime_api.h>
+#include <cuda_fp16.h>
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define FMHA_CHECK_CUDA( call )                                                                    \
+    do {                                                                                           \
+        cudaError_t status_ = call;                                                                \
+        if( status_ != cudaSuccess ) {                                                             \
+            fprintf( stderr,                                                                       \
+                     "CUDA error (%s:%d): %s\n",                                                   \
+                     __FILE__,                                                                     \
+                     __LINE__,                                                                     \
+                     cudaGetErrorString( status_ ) );                                              \
+            exit( 1 );                                                                             \
+        }                                                                                          \
+    } while( 0 )
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+enum Data_type { DATA_TYPE_FP16, DATA_TYPE_FP32, DATA_TYPE_INT32, DATA_TYPE_INT8 };
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline void set_alpha( uint32_t &alpha, float norm, Data_type dtype ) {
+    if( dtype == DATA_TYPE_FP16 ) {
+        half x = __float2half_rn( norm );
+        uint16_t h = reinterpret_cast<const uint16_t &>( x );
+        ushort2 h2 = { h, h };
+        alpha = reinterpret_cast<const uint32_t &>( h2 );
+    } else if( dtype == DATA_TYPE_FP32 ) {
+        alpha = reinterpret_cast<const uint32_t &>( norm );
+    } else if( dtype == DATA_TYPE_INT32 ) {
+        int32_t inorm = static_cast<int32_t>( norm );
+        alpha = reinterpret_cast<const uint32_t &>( inorm );
+    } else {
+        assert( false );
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline size_t get_size_in_bytes( size_t n, Data_type dtype ) {
+    switch( dtype ) {
+    case DATA_TYPE_FP32:
+        return n * 4;
+    case DATA_TYPE_FP16:
+        return n * 2;
+    case DATA_TYPE_INT32:
+        return n * 4;
+    case DATA_TYPE_INT8:
+        return n;
+    default:
+        assert( false );
+        return 0;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm.cu
new file mode 100644
index 0000000000..1ec98eeea6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm.cu
@@ -0,0 +1,328 @@
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDACachingAllocator.h>
+
+#include "batch_norm.h"
+
+#include <cuda.h>
+
+#include "compat.h"
+
+#define cudaCheckErrors(msg) \
+    do { \
+        cudaError_t __err = cudaGetLastError(); \
+        if (__err != cudaSuccess) { \
+            fprintf(stderr, "Fatal error: %s (%s at %s:%d)\n", \
+                msg, cudaGetErrorString(__err), \
+                __FILE__, __LINE__); \
+            fprintf(stderr, "*** FAILED - ABORTING\n"); \
+            exit(1); \
+        } \
+    } while (0)
+
+static size_t round_up_to_multiple(size_t x, int multiple) {
+  return ((x + multiple - 1) / multiple) * multiple;
+}
+
+struct Workspace {
+  Workspace(size_t size) : size(size), data(NULL) {
+    auto& allocator = *::c10::cuda::CUDACachingAllocator::get();
+    dataPtr = allocator.allocate(size);
+    data = dataPtr.get();
+  }
+  Workspace(const Workspace&) = delete;
+  Workspace(Workspace&&) = default;
+  Workspace& operator=(Workspace&&) = default;
+  ~Workspace() = default;
+
+  size_t size;
+  void* data;
+  c10::DataPtr dataPtr;
+};
+
+// Return {y}
+at::Tensor nhwc_bn_fwd_train(
+                       const at::Tensor& x,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& minibatch_mean,
+                       const at::Tensor& minibatch_inv_var,
+                       const at::Tensor& ret_cta,
+                       const float momentum,
+                       const float epsilon,
+                       const bool fuse_relu,
+                       void * my_data,
+                       void * pair_data,
+                       void * pair_data2,
+                       void * pair_data3,
+                       const int bn_group,
+                       const at::Tensor& magic_tensor,
+                       const int occupancy,
+                       const int grid_dim_x,
+                       const bool coop) {
+
+  const int N = x.size(0);
+  const int H = x.size(1);
+  const int W = x.size(2);
+  const int C = x.size(3);
+
+  // generating new magic number and use that for sync
+  int* magic = magic_tensor.DATA_PTR<int>();
+  *magic = (*magic + 1) & 0xff;
+
+  // Allocate output tensor
+  at::Tensor y = at::empty({N, H, W, C}, x.options());
+
+  // Create wrapper
+  NhwcBatchNorm *bn = new NhwcBatchNorm();
+
+  bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);
+  bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);
+
+  bn->setConstants(momentum, epsilon);
+
+  // set pointers within the wrapper
+  bn->setInputOutputPointers(x.DATA_PTR<at::Half>(),
+                             nullptr,
+                             y.DATA_PTR<at::Half>(),
+                             nullptr);
+
+  bn->setWeightPointers({scale.DATA_PTR<float>(), bias.DATA_PTR<float>()}, {nullptr, nullptr});
+  bn->setParameterPointers({running_mean.DATA_PTR<float>(), running_inv_var.DATA_PTR<float>()});
+
+  // deal with workspace(s)
+  auto workspace_bytes = bn->numWorkspaceBytes();
+  // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset
+  // an allocated workspace for the others
+  size_t total_workspace_bytes = 0;
+  std::vector<size_t> workspace_offsets;
+
+  for (auto index = 3; index < workspace_bytes.size(); ++index) {
+    total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);
+    workspace_offsets.push_back(total_workspace_bytes);
+
+    auto alloc_bytes = workspace_bytes[index];
+    total_workspace_bytes += alloc_bytes;
+  }
+
+  // Allocate the workspace
+  Workspace ws(total_workspace_bytes);
+
+  std::vector<void *> workspace;
+  workspace.push_back(minibatch_mean.DATA_PTR<float>());
+  workspace.push_back(minibatch_inv_var.DATA_PTR<float>());
+
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const int retired_cta_bytes = workspace_bytes[2];
+  void* retired_ctas = ret_cta.DATA_PTR<uint8_t>();
+  assert(ret_cta.size(0)>=retired_cta_bytes);
+  workspace.push_back(retired_ctas);
+
+  for (auto index = 3; index < workspace_bytes.size(); ++index) {
+    void *ptr = reinterpret_cast<uint8_t*>(ws.data) + workspace_offsets[index-3];
+    workspace.push_back(ptr);
+  }
+
+  bn->setWorkspacePointers(workspace, workspace_bytes);
+
+  // Don't fuse in ReLU for now at least
+  bn->fwd(stream, fuse_relu, my_data, pair_data, pair_data2, pair_data3, bn_group, *magic, occupancy, grid_dim_x, coop);
+
+  return y;
+}
+
+at::Tensor nhwc_bn_fwd_eval(
+                       const at::Tensor& x,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& ret_cta,
+                       const int bn_group,
+                       const float momentum,
+                       const float epsilon,
+                       const bool fuse_relu) {
+
+  const int N = x.size(0);
+  const int H = x.size(1);
+  const int W = x.size(2);
+  const int C = x.size(3);
+
+  // Allocate output tensor
+  at::Tensor y = at::empty({N, H, W, C}, x.options());
+
+  // Create wrapper
+  NhwcBatchNorm *bn = new NhwcBatchNorm();
+
+  bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);
+  bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);
+
+  bn->setConstants(momentum, epsilon);
+
+  // set pointers within the wrapper
+  bn->setInputOutputPointers(x.DATA_PTR<at::Half>(),
+                             nullptr,
+                             y.DATA_PTR<at::Half>(),
+                             nullptr);
+
+  bn->setWeightPointers({scale.DATA_PTR<float>(), bias.DATA_PTR<float>()}, {nullptr, nullptr});
+  bn->setParameterPointers({running_mean.DATA_PTR<float>(), running_inv_var.DATA_PTR<float>()});
+
+  // deal with workspace(s)
+  auto workspace_bytes = bn->numWorkspaceBytes();
+  // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset
+  // an allocated workspace for the others
+  size_t total_workspace_bytes = 0;
+  std::vector<size_t> workspace_offsets;
+
+  for (auto index = 3; index < workspace_bytes.size(); ++index) {
+    total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);
+    workspace_offsets.push_back(total_workspace_bytes);
+
+    auto alloc_bytes = workspace_bytes[index];
+    total_workspace_bytes += alloc_bytes;
+  }
+
+  // Allocate the workspace
+  Workspace ws(total_workspace_bytes);
+
+  std::vector<void *> workspace;
+  workspace.push_back(nullptr);
+  workspace.push_back(nullptr);
+
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const int retired_cta_bytes = workspace_bytes[2];
+  void* retired_ctas = ret_cta.DATA_PTR<uint8_t>();
+  assert(ret_cta.size(0)>=retired_cta_bytes);
+  workspace.push_back(retired_ctas);
+
+  for (auto index = 3; index < workspace_bytes.size(); ++index) {
+    void *ptr = reinterpret_cast<uint8_t*>(ws.data) + workspace_offsets[index-3];
+    workspace.push_back(ptr);
+  }
+
+  bn->setWorkspacePointers(workspace, workspace_bytes);
+
+  // Don't fuse in ReLU for now at least
+  bn->fwdInference(stream, fuse_relu);
+
+  return y;
+
+}
+
+std::vector<at::Tensor> nhwc_bn_bwd(
+                       const at::Tensor& x,
+                       const at::Tensor& dy,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& minibatch_mean,
+                       const at::Tensor& minibatch_inv_var,
+                       const at::Tensor& ret_cta,
+                       const float momentum,
+                       const float epsilon,
+                       const bool fuse_relu,
+                       void * my_data,
+                       void * pair_data, 
+                       void * pair_data2, 
+                       void * pair_data3, 
+                       const int bn_group,
+                       const at::Tensor& magic_tensor,
+                       const int occupancy,
+                       const int grid_dim_x,
+                       const bool coop) {
+  // shape
+  const int N = x.size(0);
+  const int H = x.size(1);
+  const int W = x.size(2);
+  const int C = x.size(3);
+
+  // generating new magic number and use that for sync
+  int* magic = magic_tensor.DATA_PTR<int>();
+  *magic = (*magic + 1) & 0xff;
+
+  // outputs
+  at::Tensor x_grad, scale_grad, bias_grad;
+
+  // Allocate outputs
+  x_grad = at::empty_like(x);
+  scale_grad = at::empty_like(scale);
+  bias_grad = at::empty_like(bias);
+
+  // Create wrapper
+  NhwcBatchNorm *bn = new NhwcBatchNorm();
+
+  bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);
+  bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);
+
+  bn->setConstants(momentum, epsilon);
+
+  // set pointers within the wrapper
+  bn->setInputOutputPointers(x.DATA_PTR<at::Half>(),
+                             x_grad.DATA_PTR<at::Half>(),
+                             nullptr,
+                             dy.DATA_PTR<at::Half>());
+
+  bn->setWeightPointers({scale.DATA_PTR<float>(), bias.DATA_PTR<float>()}, {scale_grad.DATA_PTR<float>(), bias_grad.DATA_PTR<float>()});
+  bn->setParameterPointers({running_mean.DATA_PTR<float>(), running_inv_var.DATA_PTR<float>()});
+
+  // deal with workspace(s)
+  auto workspace_bytes = bn->numWorkspaceBytes();
+  // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset
+  // an allocated workspace for the others
+  size_t total_workspace_bytes = 0;
+  std::vector<size_t> workspace_offsets;
+
+  for (auto index = 3; index < workspace_bytes.size(); ++index) {
+    total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);
+    workspace_offsets.push_back(total_workspace_bytes);
+
+    auto alloc_bytes = workspace_bytes[index];
+    total_workspace_bytes += alloc_bytes;
+  }
+
+  // Allocate the workspace
+  Workspace ws(total_workspace_bytes);
+
+  std::vector<void *> workspace;
+  workspace.push_back(minibatch_mean.DATA_PTR<float>());
+  workspace.push_back(minibatch_inv_var.DATA_PTR<float>());
+
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const int retired_cta_bytes = workspace_bytes[2];
+  void* retired_ctas = ret_cta.DATA_PTR<uint8_t>();
+  assert(ret_cta.size(0)>=retired_cta_bytes);
+  workspace.push_back(retired_ctas);
+
+  for (auto index = 3; index < workspace_bytes.size(); ++index) {
+    void *ptr = reinterpret_cast<uint8_t*>(ws.data) + workspace_offsets[index-3];
+    workspace.push_back(ptr);
+  }
+
+  bn->setWorkspacePointers(workspace, workspace_bytes);
+
+  bn->dgrad(stream, fuse_relu, my_data, pair_data, pair_data2, pair_data3, bn_group, *magic, occupancy, grid_dim_x, coop);
+
+  return std::vector<at::Tensor>{x_grad, scale_grad, bias_grad};
+}
+
+int nhwc_bn_fwd_occupancy() {
+    int device_id=-1;
+    cudaGetDevice(&device_id);
+
+    //max occupancy supported by the code is 2
+    return NhwcBatchNorm::smem_driven_fwd_occupancy(device_id, 2);
+}
+
+int nhwc_bn_bwd_occupancy() {
+    int device_id=-1;
+    cudaGetDevice(&device_id);
+    
+    //max occupancy supported by the code is 2
+    return NhwcBatchNorm::smem_driven_bwd_occupancy(device_id, 2);
+}
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm.h
new file mode 100644
index 0000000000..bb79d67589
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm.h
@@ -0,0 +1,735 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+/*!
+ * Copyright (c) 2018 by Contributors
+ * \file nhwc_batch_norm.h
+ * \brief CUDA NHWC Batch Normalization code
+ * \author Shankara Rao Thejaswi Nanditale, Dick Carter, Evgeni Krimer
+*/
+#ifndef MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_H_
+#define MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_H_
+
+#include <cudnn.h>
+
+#include <algorithm>
+#include <vector>
+#include <string>
+#include <iostream>
+
+#include "nhwc_batch_norm_kernel.h"
+#include "cuda_utils.h"
+
+
+#define VERBOSE_DEFAULT false
+
+class NhwcBatchNorm {
+ public:
+  NhwcBatchNorm() {
+    name_ = "nhwc_batchnorm";
+    createTensorDescriptor(&X_tensor_desc_);
+    createTensorDescriptor(&Y_tensor_desc_);
+  }
+
+  ~NhwcBatchNorm() {
+    destroyTensorDescriptor(X_tensor_desc_);
+    destroyTensorDescriptor(Y_tensor_desc_);
+  }
+
+  void die() {
+    std::cerr << "batchnorm not initialized" << std::endl;
+    exit(-1);
+  }
+
+  void fwd(cudaStream_t stream, bool use_relu, void* my_data, void* pair_data, void* pair_data2, void* pair_data3, const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop);
+  void dgrad(cudaStream_t stream, bool use_relu, void* my_data, void* pair_data, void* pair_data2, void* pair_data3, const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop);
+  void fwdInference(cudaStream_t stream, bool use_relu);
+  dim3 calc_fwd_grid(int *loop, const int grid_dim_x);
+  dim3 calc_bwd_grid(int *loop, const int grid_dim_x);
+
+  void setInputDescriptor(const cudnnTensorFormat_t format,
+                                  const cudnnDataType_t     data_type,
+                                  int n, int c, int h, int w, int bn_group) {
+    m_ = n * h * w;
+    int m_bn_adjusted = m_ * bn_group;
+    c_ = c;
+    // factor to scale sum of squared errors to get saved variance.  Must be 1/nhw.
+    svar_inv_count_ = 1.f / m_bn_adjusted;
+    // factor to scale sum of squared errors to get running variance. Should be 1/(nhw-1).
+    int divisor = m_bn_adjusted - 1;
+    // nhw == 1 is unlikely, but by setting the rvar_inv_count_ == 1.f, we avoid running var infs.
+    rvar_inv_count_ = divisor == 0 ? 1.f : 1.f / divisor;
+    setTensorDescriptor(X_tensor_desc_, format, data_type, n, c, h, w);
+  }
+
+  void setOutputDescriptor(const cudnnTensorFormat_t format,
+                                   const cudnnDataType_t     data_type,
+                                   int n, int c, int h, int w) {
+    setTensorDescriptor(Y_tensor_desc_, format, data_type, n, c, h, w);
+  }
+
+  const std::vector<size_t> numWorkspaceBytes() const;
+
+  void setWorkspacePointers(
+      const std::vector<void*>&  workspace,
+      const std::vector<size_t>& num_workspace_bytes);
+
+  void setInputOutputPointers(void* X, void* dX, void* Y, void *dY) {
+    X_ = X;
+    dX_  = dX;
+    Y_   = Y;
+    dY_  = dY;
+  }
+
+  // Sets the pointers for the scale and weight (in that order) data and derivative buffers.
+  void setWeightPointers(const std::vector<void*>& weight_pointers,
+                                 const std::vector<void*>& deriv_pointers) {
+    assert(weight_pointers.size() == 2);
+    assert(deriv_pointers.size()  == 2);
+    scale_  = static_cast<float*>(weight_pointers[0]);
+    bias_   = static_cast<float*>(weight_pointers[1]);
+    dscale_ = static_cast<float*>(deriv_pointers[0]);
+    dbias_  = static_cast<float*>(deriv_pointers[1]);
+  }
+
+  // Sets the pointers for the population mean and variance buffers, in that order.
+  void setParameterPointers(const std::vector<void*>& param_pointers) {
+    assert(param_pointers.size() == 2);
+    population_mean_     = static_cast<float*>(param_pointers[0]);
+    population_variance_ = static_cast<float*>(param_pointers[1]);
+  }
+
+  void setConstants(const double exp_avg_factor, const double eps) {
+    exp_avg_factor_ = exp_avg_factor;
+    eps_ = eps;
+  }
+
+  void processCudnnStatus(const cudnnStatus_t& status,
+                          const std::string& string = std::string(),
+                          bool verbose = VERBOSE_DEFAULT) {
+    if (status != CUDNN_STATUS_SUCCESS)
+      LOG(FATAL) << string << " " << cudnnGetErrorString(status);
+    else if (verbose)
+      LOG(INFO) << string << " " << cudnnGetErrorString(status);
+  }
+
+  void checkCudaStatus(const std::string& string = std::string(),
+                       bool verbose = VERBOSE_DEFAULT) {
+    cudaError_t status = cudaGetLastError();
+    if (status != cudaSuccess)
+      LOG(FATAL) << string << " " << cudaGetErrorString(status);
+    else if (verbose)
+      LOG(INFO) << string << " " << cudaGetErrorString(status);
+  }
+
+  size_t size_retired_ctas(int grid_y) const {
+    // Note that the value of max_grid_y to handle known GPUs is about 160.
+    const int max_grid_y = 1024;
+    if (grid_y > max_grid_y)
+      LOG(INFO) << "GPU capabilities exceeds assumptions.";
+    const int retired_cta_bytes = max_grid_y * 2 * sizeof(int);
+    // Since the region will be initialized once and used for many kernels,
+    // the idea is to return an ample size that will cover all uses.
+    return retired_cta_bytes;
+  }
+
+  cudnnTensorDescriptor_t  X_tensor_desc_ = nullptr;
+  cudnnTensorDescriptor_t  Y_tensor_desc_ = nullptr;
+
+  void*  X_ = nullptr;
+  void* dX_ = nullptr;
+  void*  Y_ = nullptr;
+  void* dY_ = nullptr;
+
+  // Learned scale and bias weights.
+  float* scale_  = nullptr;
+  float* dscale_ = nullptr;
+  float* bias_   = nullptr;
+  float* dbias_  = nullptr;
+
+  // Computed population mean and variance parameters.
+  float* population_mean_     = nullptr;
+  float* population_variance_ = nullptr;
+
+  // Workspace buffers for minibatch mean and variance (computed in fwd, needed by bwd).
+  float* minibatch_mean_     = nullptr;
+  float* minibatch_variance_ = nullptr;
+
+  int m_ = 0;  // Number of values per channel that BN is normalizing.
+  int c_ = 0;  // Number of channels over which BN is normalizing.
+
+  float svar_inv_count_ = 0.f;  // factor to scale sum of squared errors to get saved variance
+  float rvar_inv_count_ = 0.f;  // factor to scale sum of squared errors to get running variance
+
+  double exp_avg_factor_ = 0.;
+  double eps_            = 0.;
+  std::string name_;
+
+ private:
+  void setTensorDescriptor(cudnnTensorDescriptor_t descriptor,
+                           cudnnTensorFormat_t format,
+                           cudnnDataType_t     data_type,
+                           int n, int c, int h, int w) {
+    cudnnStatus_t status = CUDNN_STATUS_SUCCESS;
+    status = cudnnSetTensor4dDescriptor(descriptor, format, data_type, n, c, h, w);
+    processCudnnStatus(status, "set tensor descriptor");
+  }
+
+  void createTensorDescriptor(cudnnTensorDescriptor_t *descriptor) {
+    cudnnStatus_t status = CUDNN_STATUS_SUCCESS;
+    status = cudnnCreateTensorDescriptor(descriptor);
+    processCudnnStatus(status, "create tensor_descriptor");
+  }
+
+  void destroyTensorDescriptor(cudnnTensorDescriptor_t descriptor) {
+    cudnnStatus_t status = CUDNN_STATUS_SUCCESS;
+    status = cudnnDestroyTensorDescriptor(descriptor);
+    processCudnnStatus(status, "destroy tensor_descriptor");
+  }
+
+ protected:
+  float *partial_sums_ = nullptr;
+  int *partial_counts_ = nullptr;
+  int *retired_ctas_   = nullptr;
+
+  void _setFwdParams(NhwcBatchNormFwdParams *params) const;
+  void _setFwdInferenceParams(NhwcBatchNormFwdInferenceParams *params) const;
+  void _setBwdParams(NhwcBatchNormBwdParams *params) const;
+
+  // @todo: ability to configure these?
+  // Kernel params
+  static const int USE_ONLINE_APPROACH = 1;
+  static const int THREADS_PER_CTA = 512;
+  static const int THREADS_PER_PIXEL = 16;
+  static const int C_ELEMENTS_PER_CTA = 64;
+  static const int ELEMENTS_PER_LDG = C_ELEMENTS_PER_CTA / THREADS_PER_PIXEL;
+  static const int MAX_SMEM_WITHOUT_OPT_IN = 48 * 1024;
+
+  typedef uint16_t StorageType;
+  //typedef float StorageType;
+  // increasing this to 6 causes spills in fwd kernel!
+  static const int PIXELS_PER_THREAD_IN_REGISTERS_FWD = 5;
+  static const int PIXELS_PER_THREAD_IN_REGISTERS_BWD = 3;
+  static const int PIXELS_PER_THREAD_IN_SMEM_FWD = 10;
+  static const int PIXELS_PER_THREAD_IN_SMEM_BWD = 5;
+
+  static const int PIXELS_PER_THREAD_FWD = PIXELS_PER_THREAD_IN_REGISTERS_FWD + \
+      PIXELS_PER_THREAD_IN_SMEM_FWD;
+  static const int PIXELS_PER_THREAD_BWD = PIXELS_PER_THREAD_IN_REGISTERS_BWD + \
+      PIXELS_PER_THREAD_IN_SMEM_BWD;
+  static const int PIXELS_PER_THREAD_FWD_INFERENCE = 4;
+
+  // Derived params
+  static const size_t SMEM_SIZE_FWD = PIXELS_PER_THREAD_IN_SMEM_FWD*THREADS_PER_CTA*\
+      ELEMENTS_PER_LDG*sizeof(StorageType);
+  static const size_t SMEM_SIZE_BWD = PIXELS_PER_THREAD_IN_SMEM_BWD*THREADS_PER_CTA*\
+      ELEMENTS_PER_LDG*2*sizeof(StorageType);
+  static const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;
+  static const int PIXELS_PER_CTA_FWD = THREADS_PER_CTA/THREADS_PER_PIXEL * \
+      PIXELS_PER_THREAD_FWD;
+  static const int PIXELS_PER_CTA_BWD = THREADS_PER_CTA/THREADS_PER_PIXEL * \
+      PIXELS_PER_THREAD_BWD;
+  static const int PIXELS_PER_CTA_FWD_INFERENCE = THREADS_PER_CTA/THREADS_PER_PIXEL * \
+      PIXELS_PER_THREAD_FWD_INFERENCE;
+
+  // max grid.y in case of group bn is limited by exchange buffer size
+  static const int MAX_GBN_BLOCK_Y = 256;
+
+  // Helper function to launch the forward kernel.
+
+  // We calculate (based on smem usage) the achievable occupancy and make sure we run a kernel
+  // version that was compiled with that occupancy in its launch bounds.  This way, we avoid
+  // needless register spills.
+  void _fwdKernelLauncher(cudaStream_t stream, NhwcBatchNormFwdParams params,
+                                dim3 grid_dim, int outer_loops, bool use_relu, const int occupancy, const bool coop) {
+
+#define LAUNCH_FWD_KERNEL(OUTER_LOOPS, USE_RELU, USE_ADD_RELU, COMPILED_FOR_OCCUPANCY, COOP) \
+    do { \
+        CHECK(SMEM_SIZE_FWD <= MAX_SMEM_WITHOUT_OPT_IN) << "Nhwc batchnorm kernel smem too big."; \
+        auto fwd_func = nhwc_batch_norm_fwd< \
+                        StorageType, \
+                        THREADS_PER_CTA, \
+                        THREADS_PER_PIXEL, \
+                        PIXELS_PER_THREAD_IN_REGISTERS_FWD, \
+                        PIXELS_PER_THREAD_IN_SMEM_FWD, \
+                        ELEMENTS_PER_LDG, \
+                        USE_ONLINE_APPROACH, \
+                        OUTER_LOOPS, \
+                        USE_RELU, \
+                        USE_ADD_RELU, \
+                        COMPILED_FOR_OCCUPANCY>; \
+        if (COMPILED_FOR_OCCUPANCY > 1) { \
+            cudaFuncSetAttribute(fwd_func, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \
+            checkCudaStatus(name_ + " fwd ser coop kernel (cudaFuncSetAttribute carveout)"); \
+        } \
+        void *params_ptr = static_cast<void*>(&params); \
+        using FWD_FUNC = decltype(nhwc_batch_norm_fwd< \
+                        StorageType, \
+                        THREADS_PER_CTA, \
+                        THREADS_PER_PIXEL, \
+                        PIXELS_PER_THREAD_IN_REGISTERS_FWD, \
+                        PIXELS_PER_THREAD_IN_SMEM_FWD, \
+                        ELEMENTS_PER_LDG, \
+                        USE_ONLINE_APPROACH, \
+                        OUTER_LOOPS, \
+                        USE_RELU, \
+                        USE_ADD_RELU, \
+                        COMPILED_FOR_OCCUPANCY>); \
+        if (COOP) { \
+            cudaLaunchCooperativeKernel<FWD_FUNC>(fwd_func, \
+                grid_dim, \
+                THREADS_PER_CTA, \
+                &params_ptr, \
+                SMEM_SIZE_FWD, \
+                stream); \
+        } else { \
+            cudaLaunchKernel<FWD_FUNC>(fwd_func, \
+                grid_dim, \
+                THREADS_PER_CTA, \
+                &params_ptr, \
+                SMEM_SIZE_FWD, \
+                stream); \
+        } \
+        checkCudaStatus(name_ + " fwd ser coop kernel"); \
+    } while (0)
+
+    // Don't try for an occupancy > 2 as this will squeeze register use and create spills.
+    if (outer_loops == 1 && use_relu) {
+      if (occupancy >= 2)
+        LAUNCH_FWD_KERNEL(1, true, false, 2, coop);
+      else
+        LAUNCH_FWD_KERNEL(1, true, false, 1, coop);
+    } else if (outer_loops == 1 && !use_relu) {
+      if (occupancy >= 2)
+        LAUNCH_FWD_KERNEL(1, false, false, 2, coop);
+      else
+        LAUNCH_FWD_KERNEL(1, false, false, 1, coop);
+    } else if (use_relu) {
+      if (occupancy >= 2)
+        LAUNCH_FWD_KERNEL(0, true, false, 2, coop);
+      else
+        LAUNCH_FWD_KERNEL(0, true, false, 1, coop);
+    } else {
+      if (occupancy >= 2)
+        LAUNCH_FWD_KERNEL(0, false, false, 2, coop);
+      else
+        LAUNCH_FWD_KERNEL(0, false, false, 1, coop);
+    }
+#undef LAUNCH_FWD_KERNEL
+  }
+
+  // Helper function to launch the backward kernel.
+
+  void _bwdKernelLauncher(cudaStream_t stream, NhwcBatchNormBwdParams params,
+                                dim3 grid_dim, int outer_loops, bool use_relu, const int occupancy, const bool coop) {
+#define LAUNCH_BWD_KERNEL(OUTER_LOOPS, COMPILED_FOR_OCCUPANCY, COOP) \
+    do { \
+        CHECK(SMEM_SIZE_BWD <= MAX_SMEM_WITHOUT_OPT_IN) << "Nhwc batchnorm kernel smem too big."; \
+        auto bwd_func = nhwc_batch_norm_bwd< \
+                        StorageType, \
+                        THREADS_PER_CTA, \
+                        THREADS_PER_PIXEL, \
+                        PIXELS_PER_THREAD_IN_REGISTERS_BWD, \
+                        PIXELS_PER_THREAD_IN_SMEM_BWD, \
+                        ELEMENTS_PER_LDG, \
+                        USE_ONLINE_APPROACH, \
+                        OUTER_LOOPS, \
+                        COMPILED_FOR_OCCUPANCY>; \
+        if (COMPILED_FOR_OCCUPANCY > 1) { \
+            cudaFuncSetAttribute(bwd_func, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \
+            checkCudaStatus(name_ + " bwd coop serial kernel (cudaFuncSetAttribute carveout)"); \
+        } \
+        void *params_ptr = static_cast<void*>(&params); \
+        using BWD_FUNC = decltype(nhwc_batch_norm_bwd< \
+                        StorageType, \
+                        THREADS_PER_CTA, \
+                        THREADS_PER_PIXEL, \
+                        PIXELS_PER_THREAD_IN_REGISTERS_BWD, \
+                        PIXELS_PER_THREAD_IN_SMEM_BWD, \
+                        ELEMENTS_PER_LDG, \
+                        USE_ONLINE_APPROACH, \
+                        OUTER_LOOPS, \
+                        COMPILED_FOR_OCCUPANCY>); \
+        if (COOP) { \
+            cudaLaunchCooperativeKernel<BWD_FUNC>(bwd_func, \
+                grid_dim, \
+                THREADS_PER_CTA, \
+                &params_ptr, \
+                SMEM_SIZE_BWD, \
+                stream); \
+        } else { \
+            cudaLaunchKernel<BWD_FUNC>(bwd_func, \
+                grid_dim, \
+                THREADS_PER_CTA, \
+                &params_ptr, \
+                SMEM_SIZE_BWD, \
+                stream); \
+        } \
+        checkCudaStatus(name_ + " bwd coop serial kernel"); \
+    } while (0)
+
+#define LAUNCH_BWD_RELU_KERNEL(OUTER_LOOPS, COMPILED_FOR_OCCUPANCY, COOP) \
+    do { \
+        CHECK(SMEM_SIZE_BWD <= MAX_SMEM_WITHOUT_OPT_IN) << "Nhwc batchnorm kernel smem too big."; \
+        auto bwd_relu_func = nhwc_batch_norm_bwd_relu< \
+                        StorageType, \
+                        THREADS_PER_CTA, \
+                        THREADS_PER_PIXEL, \
+                        PIXELS_PER_THREAD_IN_REGISTERS_BWD, \
+                        PIXELS_PER_THREAD_IN_SMEM_BWD, \
+                        ELEMENTS_PER_LDG, \
+                        USE_ONLINE_APPROACH, \
+                        OUTER_LOOPS, \
+                        COMPILED_FOR_OCCUPANCY>; \
+        if (COMPILED_FOR_OCCUPANCY > 1) { \
+            cudaFuncSetAttribute(bwd_relu_func, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \
+            checkCudaStatus(name_ + " bwd-relu coop serial kernel (cudaFuncSetAttribute carveout)"); \
+        } \
+        void *params_ptr = static_cast<void*>(&params); \
+        using BWD_RELU_FUNC = decltype(nhwc_batch_norm_bwd_relu< \
+                        StorageType, \
+                        THREADS_PER_CTA, \
+                        THREADS_PER_PIXEL, \
+                        PIXELS_PER_THREAD_IN_REGISTERS_BWD, \
+                        PIXELS_PER_THREAD_IN_SMEM_BWD, \
+                        ELEMENTS_PER_LDG, \
+                        USE_ONLINE_APPROACH, \
+                        OUTER_LOOPS, \
+                        COMPILED_FOR_OCCUPANCY>); \
+        if (COOP) { \
+            cudaLaunchCooperativeKernel<BWD_RELU_FUNC>(bwd_relu_func, \
+                grid_dim, \
+                THREADS_PER_CTA, \
+                &params_ptr, \
+                SMEM_SIZE_BWD, \
+                stream); \
+        } else { \
+            cudaLaunchKernel<BWD_RELU_FUNC>(bwd_relu_func, \
+                grid_dim, \
+                THREADS_PER_CTA, \
+                &params_ptr, \
+                SMEM_SIZE_BWD, \
+                stream); \
+        } \
+        checkCudaStatus(name_ + " bwd-relu coop serial kernel"); \
+    } while (0)
+
+    // Don't try for an occupancy > 2 as this will squeeze register use and create spills.
+    if (outer_loops == 1 && use_relu) {
+      if (occupancy >= 2)
+        LAUNCH_BWD_RELU_KERNEL(1, 2, coop);
+      else
+        LAUNCH_BWD_RELU_KERNEL(1, 1, coop);
+    } else if (outer_loops == 1 && !use_relu) {
+      if (occupancy >= 2)
+        LAUNCH_BWD_KERNEL(1, 2, coop);
+      else
+        LAUNCH_BWD_KERNEL(1, 1, coop);
+    } else if (use_relu) {
+      if (occupancy >= 2)
+        LAUNCH_BWD_RELU_KERNEL(0, 2, coop);
+      else
+        LAUNCH_BWD_RELU_KERNEL(0, 1, coop);
+    } else {
+      if (occupancy >= 2)
+        LAUNCH_BWD_KERNEL(0, 2, coop);
+      else
+        LAUNCH_BWD_KERNEL(0, 1, coop);
+    }
+#undef LAUNCH_BWD_KERNEL
+  }
+
+ public:
+
+  // Calculate the expected fwd kernel occupancy, as dictated by shared memory usage.
+  static int smem_driven_fwd_occupancy(int device_id, const int max_cta_per_sm) {
+    using namespace at::cuda::utils;
+    int fwd_reduction_bytes = THREADS_PER_PIXEL*(THREADS_PER_CTA/32)*ELEMENTS_PER_LDG*sizeof(float);
+    int fwd_smem_bytes = SMEM_SIZE_FWD + fwd_reduction_bytes;
+    int occupancy = MaxSharedMemoryPerMultiprocessor(device_id) / fwd_smem_bytes;
+    return std::min(max_cta_per_sm, occupancy);
+  }
+
+  // Calculate the expected bwd kernel occupancy, as dictated by shared memory usage.
+  static int smem_driven_bwd_occupancy(int device_id, const int max_cta_per_sm) {
+    using namespace at::cuda::utils;
+    int bwd_reduction_bytes = THREADS_PER_PIXEL*(THREADS_PER_CTA/32)*ELEMENTS_PER_LDG*sizeof(float);
+    int bwd_smem_bytes = SMEM_SIZE_BWD + bwd_reduction_bytes;
+    int occupancy = MaxSharedMemoryPerMultiprocessor(device_id) / bwd_smem_bytes;
+    return std::min(max_cta_per_sm, occupancy);
+  }
+};
+
+const std::vector<size_t> NhwcBatchNorm::numWorkspaceBytes() const {
+  assert(c_ > 0);
+
+  // choose the max memory required between fwd/bwd passes
+  int grid_x_fwd = div_up(m_, PIXELS_PER_CTA_FWD);
+  int grid_x_bwd = div_up(m_, PIXELS_PER_CTA_BWD);
+  int grid_x = max(grid_x_fwd, grid_x_bwd);
+  int grid_y = div_up(c_, C_ELEMENTS_PER_CTA);
+
+  const size_t num_mean_bytes     = c_ * sizeof(float);
+  const size_t num_variance_bytes = num_mean_bytes;
+  const size_t size_sums          = grid_y*grid_x*THREADS_PER_PIXEL*\
+      ELEMENTS_PER_LDG*2*sizeof(float);
+  const size_t size_counts        = grid_y*grid_x*sizeof(int);
+
+  return {num_mean_bytes, num_variance_bytes,
+          size_retired_ctas(grid_y), size_sums, size_counts};
+}
+
+void NhwcBatchNorm::setWorkspacePointers(
+      const std::vector<void*>& workspace,
+      const std::vector<size_t>& num_workspace_bytes) {
+  assert(workspace.size() == 5);
+  assert(num_workspace_bytes.size() == 5);
+
+  minibatch_mean_     = static_cast<float*>(workspace[0]);
+  minibatch_variance_ = static_cast<float*>(workspace[1]);
+  retired_ctas_       = static_cast<int*>(workspace[2]);
+  partial_sums_       = static_cast<float*>(workspace[3]);
+  partial_counts_     = static_cast<int*>(workspace[4]);
+}
+
+void NhwcBatchNorm::_setFwdParams(NhwcBatchNormFwdParams *params) const {
+  params->gmem_src          = static_cast<uint16_t*>(X_);
+  params->gmem_dst          = static_cast<uint16_t*>(Y_);
+  params->gmem_src1         = nullptr;
+  params->gmem_bias         = bias_;
+  params->gmem_scale        = scale_;
+  params->gmem_running_mean = population_mean_;
+  params->gmem_running_var  = population_variance_;
+  params->gmem_saved_mean   = minibatch_mean_;
+  params->gmem_saved_var    = minibatch_variance_;
+  params->gmem_relu_bitmask = nullptr;
+  params->nhw               = m_;
+  params->c                 = c_;
+  params->svar_inv_count    = svar_inv_count_;
+  params->rvar_inv_count    = rvar_inv_count_;
+  params->gmem_sums         = partial_sums_;
+  params->gmem_counts       = partial_counts_;
+  params->gmem_retired_ctas = retired_ctas_;
+  params->var_eps           = eps_;
+  params->outer_loops       = 0;
+  params->exp_avg_factor    = static_cast<float>(exp_avg_factor_);
+  params->c_blks            = div_up(c_, C_ELEMENTS_PER_CTA);
+}
+
+void NhwcBatchNorm::_setFwdInferenceParams(NhwcBatchNormFwdInferenceParams
+                                                        *params) const {
+  params->gmem_src   = static_cast<uint16_t*>(X_);
+  params->gmem_dst   = static_cast<uint16_t*>(Y_);
+  params->gmem_src1  = nullptr;
+  params->gmem_bias  = bias_;
+  params->gmem_scale = scale_;
+  params->gmem_mean  = population_mean_;
+  params->gmem_var   = population_variance_;
+  params->nhw        = m_;
+  params->c          = c_;
+  params->var_eps    = eps_;
+}
+
+void NhwcBatchNorm::_setBwdParams(NhwcBatchNormBwdParams *params) const {
+  params->gmem_src          = static_cast<uint16_t*>(X_);
+  params->gmem_dy           = static_cast<uint16_t*>(dY_);
+  params->gmem_dst          = static_cast<uint16_t*>(dX_);
+  params->gmem_dst1         = nullptr;
+  params->gmem_relu_bitmask = nullptr;
+  params->gmem_dscale       = dscale_;
+  params->gmem_dbias        = dbias_;
+  params->gmem_scale        = scale_;
+  params->gmem_bias         = bias_;
+  params->gmem_saved_mean   = minibatch_mean_;
+  params->gmem_saved_var    = minibatch_variance_;
+  params->nhw               = m_;
+  params->c                 = c_;
+  params->svar_inv_count    = svar_inv_count_;
+  params->gmem_sums         = partial_sums_;
+  params->gmem_retired_ctas = retired_ctas_;
+  params->outer_loops       = 0;
+  params->c_blks            = div_up(c_, C_ELEMENTS_PER_CTA);
+}
+
+void NhwcBatchNorm::fwdInference(cudaStream_t stream, bool use_relu) {
+  bool ptrs_are_set =
+      X_tensor_desc_ != nullptr
+      && Y_tensor_desc_ != nullptr
+      && scale_ != nullptr
+      && bias_ != nullptr
+      //      && minibatch_mean_ != nullptr
+      //      && minibatch_variance_ != nullptr
+      && population_mean_ != nullptr
+      && population_variance_ != nullptr
+      && X_ != nullptr
+      //      && dX_ != nullptr
+      && Y_ != nullptr
+      //      && dY_ != nullptr
+      //      && dscale_ != nullptr
+      //      && dbias_ != nullptr
+      && partial_sums_   != nullptr
+      && partial_counts_ != nullptr;
+
+  if (!ptrs_are_set)
+    die();
+
+  dim3 grid_dim;
+  grid_dim.x = div_up(m_, PIXELS_PER_CTA_FWD_INFERENCE);
+  grid_dim.y = div_up(c_, C_ELEMENTS_PER_CTA);
+
+  // @todo: maybe just move this inside initialize routine?
+  NhwcBatchNormFwdInferenceParams params;
+  _setFwdInferenceParams(&params);
+
+  if (use_relu) {
+    nhwc_batch_norm_fwd_inference
+      <StorageType, THREADS_PER_CTA, THREADS_PER_PIXEL, ELEMENTS_PER_LDG, true, false>
+    <<<grid_dim, THREADS_PER_CTA, 0, stream>>>(params);
+    checkCudaStatus(name_ + " fwd_inference-relu kernel");
+  } else {
+    nhwc_batch_norm_fwd_inference
+      <StorageType, THREADS_PER_CTA, THREADS_PER_PIXEL, ELEMENTS_PER_LDG, false, false>
+    <<<grid_dim, THREADS_PER_CTA, 0, stream>>>(params);
+    checkCudaStatus(name_ + " fwd_inference kernel");
+  }
+}
+
+dim3 NhwcBatchNorm::calc_fwd_grid(int *loop, const int grid_dim_x) {
+  dim3 grid_dim;
+  grid_dim.x = div_up(m_, PIXELS_PER_CTA_FWD);
+  int c_blks = div_up(c_, C_ELEMENTS_PER_CTA);
+  unsigned int max_grid_x = grid_dim_x;
+  if (grid_dim.x <= max_grid_x) {
+    *loop = 1;
+    if (max_grid_x / grid_dim.x > 1) {
+      grid_dim.y = std::min(c_blks, static_cast<int>(max_grid_x / grid_dim.x));
+      assert(grid_dim.y<MAX_GBN_BLOCK_Y); //FIXME: turn into a loop
+    } else {
+      grid_dim.y = 1;
+    }
+  } else {
+    grid_dim.x = max_grid_x;
+    grid_dim.y = 1;
+    int nhw_in_regs = m_ - PIXELS_PER_THREAD_IN_SMEM_FWD*PIXELS_PER_LDG*grid_dim.x;
+    int pixels_per_iteration = PIXELS_PER_THREAD_IN_REGISTERS_FWD*PIXELS_PER_LDG*grid_dim.x;
+    *loop = div_up(nhw_in_regs, pixels_per_iteration);
+  }
+  return grid_dim;
+}
+
+dim3 NhwcBatchNorm::calc_bwd_grid(int *loop, const int grid_dim_x) {
+  dim3 grid_dim;
+  grid_dim.x = div_up(m_, PIXELS_PER_CTA_BWD);
+  int c_blks = div_up(c_, C_ELEMENTS_PER_CTA);
+  unsigned int max_grid_x = grid_dim_x;
+  if (grid_dim.x <= max_grid_x) {
+    *loop = 1;
+    if (max_grid_x / grid_dim.x > 1) {
+      grid_dim.y = std::min(c_blks, static_cast<int>(max_grid_x / grid_dim.x));
+      assert(grid_dim.y<MAX_GBN_BLOCK_Y); //FIXME: turn into a loop
+    } else {
+      grid_dim.y = 1;
+    }
+  } else {
+    grid_dim.x = max_grid_x;
+    grid_dim.y = 1;
+    int nhw_in_regs = m_ - PIXELS_PER_THREAD_IN_SMEM_BWD*PIXELS_PER_LDG*grid_dim.x;
+    int pixels_per_iteration = PIXELS_PER_THREAD_IN_REGISTERS_BWD*PIXELS_PER_LDG*grid_dim.x;
+    *loop = div_up(nhw_in_regs, pixels_per_iteration);
+  }
+  return grid_dim;
+}
+
+void NhwcBatchNorm::fwd(cudaStream_t stream, bool use_relu, void* my_data, void* pair_data, void* pair_data2, void* pair_data3,
+                        const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop) {
+  bool ptrs_are_set =
+      X_tensor_desc_ != nullptr
+      && Y_tensor_desc_ != nullptr
+      && scale_ != nullptr
+      && bias_ != nullptr
+      && minibatch_mean_ != nullptr
+      && minibatch_variance_ != nullptr
+      && population_mean_ != nullptr
+      && population_variance_ != nullptr
+      && X_ != nullptr
+      //      && dX_ != nullptr
+      && Y_ != nullptr
+      //      && dY_ != nullptr
+      //      && dscale_ != nullptr
+      //      && dbias_ != nullptr
+      && partial_sums_   != nullptr
+      && partial_counts_ != nullptr
+      && retired_ctas_   != nullptr;
+
+  if (!ptrs_are_set)
+    die();
+
+  // reset of retired_cta_count no longer needed
+
+  NhwcBatchNormFwdParams params;
+  _setFwdParams(&params);
+  params.my_data = my_data;
+  params.pair_datas[0] = pair_data;
+  params.pair_datas[1] = pair_data2;
+  params.pair_datas[2] = pair_data3;
+  params.magic = magic;
+  params.sync_iters = (bn_group==8)?3:(bn_group >> 1);
+
+  dim3 grid_dim = calc_fwd_grid(&params.outer_loops, grid_dim_x);
+  _fwdKernelLauncher(stream, params, grid_dim, params.outer_loops, use_relu, occupancy, coop);
+}
+
+void NhwcBatchNorm::dgrad(cudaStream_t stream, bool use_relu, void* my_data, void* pair_data, void* pair_data2, void* pair_data3, 
+                          const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop) {
+  bool ptrs_are_set =
+      X_tensor_desc_ != nullptr
+      && Y_tensor_desc_ != nullptr
+      && scale_ != nullptr
+      && (bias_ != nullptr || !use_relu)
+      && minibatch_mean_ != nullptr
+      && minibatch_variance_ != nullptr
+      //      && population_mean_ != nullptr
+      //      && population_variance_ != nullptr
+      && X_ != nullptr
+      && dX_ != nullptr
+      //      && Y_ != nullptr
+      && dY_ != nullptr
+      && dscale_ != nullptr
+      && dbias_ != nullptr;
+
+  if (!ptrs_are_set)
+    die();
+
+  // reset of retired_cta_count no longer needed
+
+  NhwcBatchNormBwdParams params;
+  _setBwdParams(&params);
+  params.my_data = my_data;
+  params.pair_datas[0] = pair_data;
+  params.pair_datas[1] = pair_data2;
+  params.pair_datas[2] = pair_data3;
+  params.magic = magic;
+  params.sync_iters = (bn_group==8)?3:(bn_group >> 1);
+  params.wgrad_coeff = 1.0 / bn_group;
+
+  dim3 grid_dim = calc_bwd_grid(&params.outer_loops, grid_dim_x);
+  _bwdKernelLauncher(stream, params, grid_dim, params.outer_loops, use_relu, occupancy, coop);
+}
+
+#endif  // MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_H_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm_add_relu.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm_add_relu.cu
new file mode 100644
index 0000000000..e383fb800f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm_add_relu.cu
@@ -0,0 +1,340 @@
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <c10/cuda/CUDACachingAllocator.h>
+
+#include "batch_norm_add_relu.h"
+
+#include <cuda.h>
+
+#include "compat.h"
+
+//FIXME move the common stuff to common h file
+#define cudaCheckErrors(msg) \
+    do { \
+        cudaError_t __err = cudaGetLastError(); \
+        if (__err != cudaSuccess) { \
+            fprintf(stderr, "Fatal error: %s (%s at %s:%d)\n", \
+                msg, cudaGetErrorString(__err), \
+                __FILE__, __LINE__); \
+            fprintf(stderr, "*** FAILED - ABORTING\n"); \
+            exit(1); \
+        } \
+    } while (0)
+
+static size_t round_up_to_multiple(size_t x, int multiple) {
+  return ((x + multiple - 1) / multiple) * multiple;
+}
+
+struct Workspace {
+  Workspace(size_t size) : size(size), data(NULL) {
+    auto& allocator = *::c10::cuda::CUDACachingAllocator::get();
+    dataPtr = allocator.allocate(size);
+    data = dataPtr.get();
+  }
+  Workspace(const Workspace&) = delete;
+  Workspace(Workspace&&) = default;
+  Workspace& operator=(Workspace&&) = default;
+  ~Workspace() = default;
+
+  size_t size;
+  void* data;
+  c10::DataPtr dataPtr;
+};
+
+// Return {y}
+at::Tensor nhwc_bn_addrelu_fwd_train(
+                       const at::Tensor& x,
+                       const at::Tensor& z,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& minibatch_mean,
+                       const at::Tensor& minibatch_inv_var,
+                       const at::Tensor& bitmask,
+                       const at::Tensor& ret_cta,
+                       const float momentum,
+                       const float epsilon,
+                       void * my_data,
+                       void * pair_data,
+                       void * pair_data2,
+                       void * pair_data3,
+                       const int bn_group,
+                       const at::Tensor& magic_tensor,
+                       const int occupancy,
+                       const int grid_dim_x,
+                       const bool coop) {
+
+  const int N = x.size(0);
+  const int H = x.size(1);
+  const int W = x.size(2);
+  const int C = x.size(3);
+
+  // generating new magic number and use that for sync
+  int* magic = magic_tensor.DATA_PTR<int>();
+  *magic = (*magic + 1) & 0xff;
+
+  // Allocate output tensor
+  at::Tensor y = at::empty({N, H, W, C}, x.options());
+
+  // Create wrapper
+  NhwcBatchNormAddRelu *bn = new NhwcBatchNormAddRelu();
+
+  bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);
+  bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);
+
+  bn->setConstants(momentum, epsilon);
+
+  // set pointers within the wrapper
+  bn->setInputOutputPointers(x.DATA_PTR<at::Half>(),
+                             nullptr,
+                             y.DATA_PTR<at::Half>(),
+                             nullptr,
+                             z.DATA_PTR<at::Half>(),
+                             nullptr);
+
+  bn->setWeightPointers({scale.DATA_PTR<float>(), bias.DATA_PTR<float>()}, {nullptr, nullptr});
+  bn->setParameterPointers({running_mean.DATA_PTR<float>(), running_inv_var.DATA_PTR<float>()});
+
+  // deal with workspace(s)
+  auto workspace_bytes = bn->numWorkspaceBytes();
+  // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset
+  // an allocated workspace for the others
+  size_t total_workspace_bytes = 0;
+  std::vector<size_t> workspace_offsets;
+
+  for (auto index = 4; index < workspace_bytes.size(); ++index) {
+    total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);
+    workspace_offsets.push_back(total_workspace_bytes);
+
+    auto alloc_bytes = workspace_bytes[index];
+    total_workspace_bytes += alloc_bytes;
+  }
+
+  // Allocate the workspace
+  Workspace ws(total_workspace_bytes);
+
+  std::vector<void *> workspace;
+  workspace.push_back(minibatch_mean.DATA_PTR<float>());
+  workspace.push_back(minibatch_inv_var.DATA_PTR<float>());
+  workspace.push_back(bitmask.DATA_PTR<int32_t>());
+
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const int retired_cta_bytes = workspace_bytes[3];
+  void* retired_ctas = ret_cta.DATA_PTR<uint8_t>();
+  assert(ret_cta.size(0)>=retired_cta_bytes);
+
+  workspace.push_back(retired_ctas);
+
+  for (auto index = 4; index < workspace_bytes.size(); ++index) {
+    void *ptr = reinterpret_cast<uint8_t*>(ws.data) + workspace_offsets[index-4];
+    workspace.push_back(ptr);
+  }
+
+  bn->setWorkspacePointers(workspace, workspace_bytes);
+
+  // Don't fuse in ReLU for now at least
+  bn->fwd(stream, my_data, pair_data, pair_data2, pair_data3, bn_group, *magic, occupancy, grid_dim_x, coop);
+
+  return y;
+}
+
+at::Tensor nhwc_bn_addrelu_fwd_eval(
+                       const at::Tensor& x,
+                       const at::Tensor& z,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& ret_cta,
+                       const int bn_group,
+                       const float momentum,
+                       const float epsilon) {
+
+  const int N = x.size(0);
+  const int H = x.size(1);
+  const int W = x.size(2);
+  const int C = x.size(3);
+
+  // Allocate output tensor
+  at::Tensor y = at::empty({N, H, W, C}, x.options());
+
+  // Create wrapper
+  NhwcBatchNormAddRelu *bn = new NhwcBatchNormAddRelu();
+
+  bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);
+  bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);
+
+  bn->setConstants(momentum, epsilon);
+
+  // set pointers within the wrapper
+  bn->setInputOutputPointers(x.DATA_PTR<at::Half>(),
+                             nullptr,
+                             y.DATA_PTR<at::Half>(),
+                             nullptr,
+                             z.DATA_PTR<at::Half>(),
+                             nullptr);
+
+  bn->setWeightPointers({scale.DATA_PTR<float>(), bias.DATA_PTR<float>()}, {nullptr, nullptr});
+  bn->setParameterPointers({running_mean.DATA_PTR<float>(), running_inv_var.DATA_PTR<float>()});
+
+  // deal with workspace(s)
+  auto workspace_bytes = bn->numWorkspaceBytes();
+  // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset
+  // an allocated workspace for the others
+  size_t total_workspace_bytes = 0;
+  std::vector<size_t> workspace_offsets;
+
+  for (auto index = 4; index < workspace_bytes.size(); ++index) {
+    total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);
+    workspace_offsets.push_back(total_workspace_bytes);
+
+    auto alloc_bytes = workspace_bytes[index];
+    total_workspace_bytes += alloc_bytes;
+  }
+
+  // Allocate the workspace
+  Workspace ws(total_workspace_bytes);
+
+  std::vector<void *> workspace;
+  workspace.push_back(nullptr);
+  workspace.push_back(nullptr);
+  workspace.push_back(nullptr);
+
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const int retired_cta_bytes = workspace_bytes[3];
+  void* retired_ctas = ret_cta.DATA_PTR<uint8_t>();
+  assert(ret_cta.size(0)>=retired_cta_bytes);
+  workspace.push_back(retired_ctas);
+
+  for (auto index = 4; index < workspace_bytes.size(); ++index) {
+    void *ptr = reinterpret_cast<uint8_t*>(ws.data) + workspace_offsets[index-4];
+    workspace.push_back(ptr);
+  }
+
+  bn->setWorkspacePointers(workspace, workspace_bytes);
+
+  // Don't fuse in ReLU for now at least
+  bn->fwdInference(stream);
+
+  return y;
+
+}
+
+std::vector<at::Tensor> nhwc_bn_addrelu_bwd(
+                       const at::Tensor& x,
+                       const at::Tensor& dy,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& minibatch_mean,
+                       const at::Tensor& minibatch_inv_var,
+                       const at::Tensor& bitmask,
+                       const at::Tensor& ret_cta,
+                       const float momentum,
+                       const float epsilon,
+                       void * my_data,
+                       void * pair_data, 
+                       void * pair_data2, 
+                       void * pair_data3, 
+                       const int bn_group,
+                       const at::Tensor& magic_tensor,
+                       const int occupancy,
+                       const int grid_dim_x,
+                       const bool coop) {
+  // shape
+  const int N = x.size(0);
+  const int H = x.size(1);
+  const int W = x.size(2);
+  const int C = x.size(3);
+
+  // generating new magic number and use that for sync
+  int* magic = magic_tensor.DATA_PTR<int>();
+  *magic = (*magic + 1) & 0xff;
+
+  // outputs
+  at::Tensor x_grad, z_grad, scale_grad, bias_grad;
+
+  // Allocate outputs
+  x_grad = at::empty_like(x);
+  z_grad = at::empty_like(x);
+  scale_grad = at::empty_like(scale);
+  bias_grad = at::empty_like(bias);
+
+  // Create wrapper
+  NhwcBatchNormAddRelu *bn = new NhwcBatchNormAddRelu();
+
+  bn->setInputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W, bn_group);
+  bn->setOutputDescriptor(CUDNN_TENSOR_NHWC, CUDNN_DATA_HALF, N, C, H, W);
+
+  bn->setConstants(momentum, epsilon);
+
+  // set pointers within the wrapper
+  bn->setInputOutputPointers(x.DATA_PTR<at::Half>(),
+                             x_grad.DATA_PTR<at::Half>(),
+                             nullptr,
+                             dy.DATA_PTR<at::Half>(),
+                             nullptr,
+                             z_grad.DATA_PTR<at::Half>());
+
+  bn->setWeightPointers({scale.DATA_PTR<float>(), bias.DATA_PTR<float>()}, {scale_grad.DATA_PTR<float>(), bias_grad.DATA_PTR<float>()});
+  bn->setParameterPointers({running_mean.DATA_PTR<float>(), running_inv_var.DATA_PTR<float>()});
+
+  // deal with workspace(s)
+  auto workspace_bytes = bn->numWorkspaceBytes();
+  // We'll create explicit tensors for the first 2 workspace ptrs, then allocate & offset
+  // an allocated workspace for the others
+  size_t total_workspace_bytes = 0;
+  std::vector<size_t> workspace_offsets;
+
+  for (auto index = 4; index < workspace_bytes.size(); ++index) {
+    total_workspace_bytes = round_up_to_multiple(total_workspace_bytes, 512);
+    workspace_offsets.push_back(total_workspace_bytes);
+
+    auto alloc_bytes = workspace_bytes[index];
+    total_workspace_bytes += alloc_bytes;
+  }
+
+  // Allocate the workspace
+  Workspace ws(total_workspace_bytes);
+
+  std::vector<void *> workspace;
+  workspace.push_back(minibatch_mean.DATA_PTR<float>());
+  workspace.push_back(minibatch_inv_var.DATA_PTR<float>());
+  workspace.push_back(bitmask.DATA_PTR<int32_t>());
+
+  auto stream = at::cuda::getCurrentCUDAStream().stream();
+  const int retired_cta_bytes = workspace_bytes[3];
+  void* retired_ctas = ret_cta.DATA_PTR<uint8_t>();
+  assert(ret_cta.size(0)>=retired_cta_bytes);
+  workspace.push_back(retired_ctas);
+
+  for (auto index = 4; index < workspace_bytes.size(); ++index) {
+    void *ptr = reinterpret_cast<uint8_t*>(ws.data) + workspace_offsets[index-4];
+    workspace.push_back(ptr);
+  }
+
+  bn->setWorkspacePointers(workspace, workspace_bytes);
+
+  bn->dgrad(stream, my_data, pair_data, pair_data2, pair_data3, bn_group, *magic, occupancy, grid_dim_x, coop);
+
+  return std::vector<at::Tensor>{x_grad, z_grad, scale_grad, bias_grad};
+}
+
+int nhwc_bn_addrelu_fwd_occupancy() {
+    int device_id=-1;
+    cudaGetDevice(&device_id);
+    
+    //max occupancy supported by the code is 2
+    return NhwcBatchNormAddRelu::smem_driven_fwd_occupancy(device_id, 2);
+}
+
+int nhwc_bn_addrelu_bwd_occupancy() {
+    int device_id=-1;
+    cudaGetDevice(&device_id);
+
+    //max occupancy supported by the code is 2
+    return NhwcBatchNormAddRelu::smem_driven_bwd_occupancy(device_id, 2);
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm_add_relu.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm_add_relu.h
new file mode 100644
index 0000000000..3dfe7b2698
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/batch_norm_add_relu.h
@@ -0,0 +1,682 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+/*!
+ * Copyright (c) 2018 by Contributors
+ * \file nhwc_batch_norm_add_relu.h
+ * \brief CUDA NHWC Batch Normalization code with fused addition
+ * \author Shankara Rao Thejaswi Nanditale, Dick Carter, Maxim Milakov, Evgeni Krimer
+*/
+#ifndef MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_ADD_RELU_H_
+#define MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_ADD_RELU_H_
+
+#include <cudnn.h>
+
+#include <algorithm>
+#include <vector>
+#include <string>
+#include <iostream>
+
+#include "nhwc_batch_norm_kernel.h"
+#include "cuda_utils.h"
+
+
+#define VERBOSE_DEFAULT false
+
+class NhwcBatchNormAddRelu {
+ public:
+  NhwcBatchNormAddRelu() {
+    name_ = "nhwc_batchnormaddrelu";
+    createTensorDescriptor(&X_tensor_desc_);
+    createTensorDescriptor(&Y_tensor_desc_);
+  }
+
+  ~NhwcBatchNormAddRelu() {
+    destroyTensorDescriptor(X_tensor_desc_);
+    destroyTensorDescriptor(Y_tensor_desc_);
+  }
+
+  void die() {
+    std::cerr << "batchnormaddrelu not initialized" << std::endl;
+    exit(-1);
+  }
+
+  void fwd(cudaStream_t stream, void* my_data, void* pair_data, void* pair_data2, void* pair_data3, const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop);
+  void dgrad(cudaStream_t stream, void* my_data, void* pair_data, void* pair_data2, void* pair_data3, const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop);
+  void fwdInference(cudaStream_t stream);
+  dim3 calc_fwd_grid(int *loop, const int grid_dim_x);
+  dim3 calc_bwd_grid(int *loop, const int grid_dim_x);
+
+  void setInputDescriptor(const cudnnTensorFormat_t format,
+                                  const cudnnDataType_t     data_type,
+                                  int n, int c, int h, int w, int bn_group) {
+    m_ = n * h * w;
+    int m_bn_adjusted = m_ * bn_group;
+    c_ = c;
+    // factor to scale sum of squared errors to get saved variance.  Must be 1/nhw.
+    svar_inv_count_ = 1.f / m_bn_adjusted;
+    // factor to scale sum of squared errors to get running variance. Should be 1/(nhw-1).
+    int divisor = m_bn_adjusted - 1;
+    // nhw == 1 is unlikely, but by setting the rvar_inv_count_ == 1.f, we avoid running var infs.
+    rvar_inv_count_ = divisor == 0 ? 1.f : 1.f / divisor;
+    setTensorDescriptor(X_tensor_desc_, format, data_type, n, c, h, w);
+  }
+
+  void setOutputDescriptor(const cudnnTensorFormat_t format,
+                                   const cudnnDataType_t     data_type,
+                                   int n, int c, int h, int w) {
+    setTensorDescriptor(Y_tensor_desc_, format, data_type, n, c, h, w);
+  }
+
+  const std::vector<size_t> numWorkspaceBytes() const;
+
+  void setWorkspacePointers(
+      const std::vector<void*>&  workspace,
+      const std::vector<size_t>& num_workspace_bytes);
+
+  void setInputOutputPointers(void* X, void* dX, void* Y, void *dY, void* addend, void* dAddend) {
+    X_ = X;
+    dX_  = dX;
+    Y_   = Y;
+    dY_  = dY;
+    addend_   = addend;
+    dAddend_  = dAddend;
+  }
+
+  // Sets the pointers for the scale and weight (in that order) data and derivative buffers.
+  void setWeightPointers(const std::vector<void*>& weight_pointers,
+                                 const std::vector<void*>& deriv_pointers) {
+    assert(weight_pointers.size() == 2);
+    assert(deriv_pointers.size()  == 2);
+    scale_  = static_cast<float*>(weight_pointers[0]);
+    bias_   = static_cast<float*>(weight_pointers[1]);
+    dscale_ = static_cast<float*>(deriv_pointers[0]);
+    dbias_  = static_cast<float*>(deriv_pointers[1]);
+  }
+
+  // Sets the pointers for the population mean and variance buffers, in that order.
+  void setParameterPointers(const std::vector<void*>& param_pointers) {
+    assert(param_pointers.size() == 2);
+    population_mean_     = static_cast<float*>(param_pointers[0]);
+    population_variance_ = static_cast<float*>(param_pointers[1]);
+  }
+
+  void setConstants(const double exp_avg_factor, const double eps) {
+    exp_avg_factor_ = exp_avg_factor;
+    eps_ = eps;
+  }
+
+  void processCudnnStatus(const cudnnStatus_t& status,
+                          const std::string& string = std::string(),
+                          bool verbose = VERBOSE_DEFAULT) {
+    if (status != CUDNN_STATUS_SUCCESS)
+      LOG(FATAL) << string << " " << cudnnGetErrorString(status);
+    else if (verbose)
+      LOG(INFO) << string << " " << cudnnGetErrorString(status);
+  }
+
+  void checkCudaStatus(const std::string& string = std::string(),
+                       bool verbose = VERBOSE_DEFAULT) {
+    cudaError_t status = cudaGetLastError();
+    if (status != cudaSuccess)
+      LOG(FATAL) << string << " " << cudaGetErrorString(status);
+    else if (verbose)
+      LOG(INFO) << string << " " << cudaGetErrorString(status);
+  }
+
+  size_t size_retired_ctas(int grid_y) const {
+    // Note that the value of max_grid_y to handle known GPUs is about 160.
+    const int max_grid_y = 1024;
+    if (grid_y > max_grid_y)
+      LOG(INFO) << "GPU capabilities exceeds assumptions.";
+    const int retired_cta_bytes = max_grid_y * 2 * sizeof(int);
+    // Since the region will be initialized once and used for many kernels,
+    // the idea is to return an ample size that will cover all uses.
+    return retired_cta_bytes;
+  }
+
+  cudnnTensorDescriptor_t  X_tensor_desc_ = nullptr;
+  cudnnTensorDescriptor_t  Y_tensor_desc_ = nullptr;
+
+  void*  X_ = nullptr;
+  void* dX_ = nullptr;
+  void*  Y_ = nullptr;
+  void* dY_ = nullptr;
+  void*  addend_ = nullptr;
+  void* dAddend_ = nullptr;
+
+  // Learned scale and bias weights.
+  float* scale_  = nullptr;
+  float* dscale_ = nullptr;
+  float* bias_   = nullptr;
+  float* dbias_  = nullptr;
+
+  // Computed population mean and variance parameters.
+  float* population_mean_     = nullptr;
+  float* population_variance_ = nullptr;
+
+  // Workspace buffers for minibatch mean and variance (computed in fwd, needed by bwd).
+  float* minibatch_mean_     = nullptr;
+  float* minibatch_variance_ = nullptr;
+
+  int m_ = 0;  // Number of values per channel that BN is normalizing.
+  int c_ = 0;  // Number of channels over which BN is normalizing.
+
+  float svar_inv_count_ = 0.f;  // factor to scale sum of squared errors to get saved variance
+  float rvar_inv_count_ = 0.f;  // factor to scale sum of squared errors to get running variance
+
+  double exp_avg_factor_ = 0.;
+  double eps_            = 0.;
+  std::string name_;
+
+ private:
+  void setTensorDescriptor(cudnnTensorDescriptor_t descriptor,
+                           cudnnTensorFormat_t format,
+                           cudnnDataType_t     data_type,
+                           int n, int c, int h, int w) {
+    cudnnStatus_t status = CUDNN_STATUS_SUCCESS;
+    status = cudnnSetTensor4dDescriptor(descriptor, format, data_type, n, c, h, w);
+    processCudnnStatus(status, "set tensor descriptor");
+  }
+
+  void createTensorDescriptor(cudnnTensorDescriptor_t *descriptor) {
+    cudnnStatus_t status = CUDNN_STATUS_SUCCESS;
+    status = cudnnCreateTensorDescriptor(descriptor);
+    processCudnnStatus(status, "create tensor_descriptor");
+  }
+
+  void destroyTensorDescriptor(cudnnTensorDescriptor_t descriptor) {
+    cudnnStatus_t status = CUDNN_STATUS_SUCCESS;
+    status = cudnnDestroyTensorDescriptor(descriptor);
+    processCudnnStatus(status, "destroy tensor_descriptor");
+  }
+
+ protected:
+  float *partial_sums_ = nullptr;
+  int *partial_counts_ = nullptr;
+  int *retired_ctas_   = nullptr;
+  unsigned int *relu_bitmask_ = nullptr;
+
+  void _setFwdParams(NhwcBatchNormFwdParams *params) const;
+  void _setFwdInferenceParams(NhwcBatchNormFwdInferenceParams *params) const;
+  void _setBwdParams(NhwcBatchNormBwdParams *params) const;
+
+  // @todo: ability to configure these?
+  // Kernel params
+  static const int USE_ONLINE_APPROACH = 1;
+  static const int THREADS_PER_CTA = 512;
+  static const int THREADS_PER_PIXEL = 16;
+  static const int C_ELEMENTS_PER_CTA = 64;
+  static const int ELEMENTS_PER_LDG = C_ELEMENTS_PER_CTA / THREADS_PER_PIXEL;
+  static const int MAX_SMEM_WITHOUT_OPT_IN = 48 * 1024;
+
+  typedef uint16_t StorageType;
+  // increasing this to 6 causes spills in fwd kernel!
+  static const int PIXELS_PER_THREAD_IN_REGISTERS_FWD = 5;
+  static const int PIXELS_PER_THREAD_IN_REGISTERS_BWD = 3;
+  static const int PIXELS_PER_THREAD_IN_SMEM_FWD = 10;
+  static const int PIXELS_PER_THREAD_IN_SMEM_BWD = 5;
+
+  static const int PIXELS_PER_THREAD_FWD = PIXELS_PER_THREAD_IN_REGISTERS_FWD + \
+      PIXELS_PER_THREAD_IN_SMEM_FWD;
+  static const int PIXELS_PER_THREAD_BWD = PIXELS_PER_THREAD_IN_REGISTERS_BWD + \
+      PIXELS_PER_THREAD_IN_SMEM_BWD;
+  static const int PIXELS_PER_THREAD_FWD_INFERENCE = 4;
+
+  // Derived params
+  static const size_t SMEM_SIZE_FWD = PIXELS_PER_THREAD_IN_SMEM_FWD*THREADS_PER_CTA*\
+      ELEMENTS_PER_LDG*sizeof(StorageType);
+  static const size_t SMEM_SIZE_BWD = PIXELS_PER_THREAD_IN_SMEM_BWD*THREADS_PER_CTA*\
+      ELEMENTS_PER_LDG*2*sizeof(StorageType);
+  static const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;
+  static const int PIXELS_PER_CTA_FWD = THREADS_PER_CTA/THREADS_PER_PIXEL * \
+      PIXELS_PER_THREAD_FWD;
+  static const int PIXELS_PER_CTA_BWD = THREADS_PER_CTA/THREADS_PER_PIXEL * \
+      PIXELS_PER_THREAD_BWD;
+  static const int PIXELS_PER_CTA_FWD_INFERENCE = THREADS_PER_CTA/THREADS_PER_PIXEL * \
+      PIXELS_PER_THREAD_FWD_INFERENCE;
+
+  // max grid.y in case of group bn is limited by exchange buffer size
+  static const int MAX_GBN_BLOCK_Y = 256;
+
+  // Helper function to launch the forward kernel.
+
+  // We calculate (based on smem usage) the achievable occupancy and make sure we run a kernel
+  // version that was compiled with that occupancy in its launch bounds.  This way, we avoid
+  // needless register spills.
+  void _fwdKernelLauncher(cudaStream_t stream, NhwcBatchNormFwdParams params,
+                                dim3 grid_dim, int outer_loops, const int occupancy, const bool coop) {
+#define LAUNCH_FWD_KERNEL(OUTER_LOOPS, USE_RELU, USE_ADD_RELU, COMPILED_FOR_OCCUPANCY, COOP) \
+    do { \
+        CHECK(SMEM_SIZE_FWD <= MAX_SMEM_WITHOUT_OPT_IN) << \
+            "Nhwc batchnormaddrelu kernel smem too big."; \
+        auto fwd_func = nhwc_batch_norm_fwd< \
+                        StorageType, \
+                        THREADS_PER_CTA, \
+                        THREADS_PER_PIXEL, \
+                        PIXELS_PER_THREAD_IN_REGISTERS_FWD, \
+                        PIXELS_PER_THREAD_IN_SMEM_FWD, \
+                        ELEMENTS_PER_LDG, \
+                        USE_ONLINE_APPROACH, \
+                        OUTER_LOOPS, \
+                        USE_RELU, \
+                        USE_ADD_RELU, \
+                        COMPILED_FOR_OCCUPANCY>; \
+        if (COMPILED_FOR_OCCUPANCY > 1) { \
+            cudaFuncSetAttribute(fwd_func, cudaFuncAttributePreferredSharedMemoryCarveout, 100); \
+            checkCudaStatus(name_ + " fwd ser coop kernel (cudaFuncSetAttribute carveout)"); \
+        } \
+        void *params_ptr = static_cast<void*>(&params); \
+        using FWD_FUNC = decltype(nhwc_batch_norm_fwd< \
+                        StorageType, \
+                        THREADS_PER_CTA, \
+                        THREADS_PER_PIXEL, \
+                        PIXELS_PER_THREAD_IN_REGISTERS_FWD, \
+                        PIXELS_PER_THREAD_IN_SMEM_FWD, \
+                        ELEMENTS_PER_LDG, \
+                        USE_ONLINE_APPROACH, \
+                        OUTER_LOOPS, \
+                        USE_RELU, \
+                        USE_ADD_RELU, \
+                        COMPILED_FOR_OCCUPANCY>); \
+        if (COOP) { \
+            cudaLaunchCooperativeKernel<FWD_FUNC>(fwd_func, \
+                grid_dim, \
+                THREADS_PER_CTA, \
+                &params_ptr, \
+                SMEM_SIZE_FWD, \
+                stream); \
+        } else { \
+            cudaLaunchKernel<FWD_FUNC>(fwd_func, \
+                grid_dim, \
+                THREADS_PER_CTA, \
+                &params_ptr, \
+                SMEM_SIZE_FWD, \
+                stream); \
+        } \
+        checkCudaStatus(name_ + " fwd ser coop kernel"); \
+    } while (0)
+
+    // Don't try for an occupancy > 2 as this will squeeze register use and create spills.
+    if (outer_loops == 1) {
+      if (occupancy >= 2)
+        LAUNCH_FWD_KERNEL(1, false, true, 2, coop);
+      else
+        LAUNCH_FWD_KERNEL(1, false, true, 1, coop);
+    } else {
+      if (occupancy >= 2)
+        LAUNCH_FWD_KERNEL(0, false, true, 2, coop);
+      else
+        LAUNCH_FWD_KERNEL(0, false, true, 1, coop);
+    }
+#undef LAUNCH_FWD_KERNEL
+  }
+
+  // Helper function to launch the backward kernel.
+
+  void _bwdKernelLauncher(cudaStream_t stream, NhwcBatchNormBwdParams params,
+                                dim3 grid_dim, int outer_loops, const int occupancy, const bool coop) {
+#define LAUNCH_BWD_ADD_RELU_KERNEL(OUTER_LOOPS, COMPILED_FOR_OCCUPANCY, COOP) \
+    do { \
+        CHECK(SMEM_SIZE_BWD <= MAX_SMEM_WITHOUT_OPT_IN) << \
+            "Nhwc batchnormaddrelu kernel smem too big."; \
+        auto bwd_add_relu_func = nhwc_batch_norm_bwd_add_relu< \
+                        StorageType, \
+                        THREADS_PER_CTA, \
+                        THREADS_PER_PIXEL, \
+                        PIXELS_PER_THREAD_IN_REGISTERS_BWD, \
+                        PIXELS_PER_THREAD_IN_SMEM_BWD, \
+                        ELEMENTS_PER_LDG, \
+                        USE_ONLINE_APPROACH, \
+                        OUTER_LOOPS, \
+                        COMPILED_FOR_OCCUPANCY>; \
+        if (COMPILED_FOR_OCCUPANCY > 1) { \
+            cudaFuncSetAttribute(bwd_add_relu_func, \
+                             cudaFuncAttributePreferredSharedMemoryCarveout, 100); \
+            checkCudaStatus(name_ + \
+                " bwd-add-relu coop serial kernel (cudaFuncSetAttribute carveout)"); \
+        } \
+        void *params_ptr = static_cast<void*>(&params); \
+        using BWD_ADD_RELU_FUNC = decltype(nhwc_batch_norm_bwd_add_relu< \
+                        StorageType, \
+                        THREADS_PER_CTA, \
+                        THREADS_PER_PIXEL, \
+                        PIXELS_PER_THREAD_IN_REGISTERS_BWD, \
+                        PIXELS_PER_THREAD_IN_SMEM_BWD, \
+                        ELEMENTS_PER_LDG, \
+                        USE_ONLINE_APPROACH, \
+                        OUTER_LOOPS, \
+                        COMPILED_FOR_OCCUPANCY>); \
+        if (COOP) { \
+            cudaLaunchCooperativeKernel<BWD_ADD_RELU_FUNC>(bwd_add_relu_func, \
+                grid_dim, \
+                THREADS_PER_CTA, \
+                &params_ptr, \
+                SMEM_SIZE_BWD, \
+                stream); \
+        } else { \
+            cudaLaunchKernel<BWD_ADD_RELU_FUNC>(bwd_add_relu_func, \
+                grid_dim, \
+                THREADS_PER_CTA, \
+                &params_ptr, \
+                SMEM_SIZE_BWD, \
+                stream); \
+        } \
+        checkCudaStatus(name_ + " bwd-add-relu coop serial kernel"); \
+  } while (0)
+
+    // Don't try for an occupancy > 2 as this will squeeze register use and create spills.
+    if (outer_loops == 1) {
+      if (occupancy >= 2)
+        LAUNCH_BWD_ADD_RELU_KERNEL(1, 2, coop);
+      else
+        LAUNCH_BWD_ADD_RELU_KERNEL(1, 1, coop);
+    } else {
+      if (occupancy >= 2)
+        LAUNCH_BWD_ADD_RELU_KERNEL(0, 2, coop);
+      else
+        LAUNCH_BWD_ADD_RELU_KERNEL(0, 1, coop);
+    }
+#undef LAUNCH_BWD_KERNEL
+  }
+
+ public:
+  // Calculate the expected fwd kernel occupancy, as dictated by shared memory usage.
+  static int smem_driven_fwd_occupancy(int device_id, const int max_cta_per_sm) {
+    using namespace at::cuda::utils;
+    int fwd_reduction_bytes = THREADS_PER_PIXEL*(THREADS_PER_CTA/32)*ELEMENTS_PER_LDG*sizeof(float);
+    int fwd_smem_bytes = SMEM_SIZE_FWD + fwd_reduction_bytes;
+    int occupancy = MaxSharedMemoryPerMultiprocessor(device_id) / fwd_smem_bytes;
+    return std::min(max_cta_per_sm, occupancy);
+  }
+
+  // Calculate the expected bwd kernel occupancy, as dictated by shared memory usage.
+  static int smem_driven_bwd_occupancy(int device_id, const int max_cta_per_sm) {
+    using namespace at::cuda::utils;
+    int bwd_reduction_bytes = THREADS_PER_PIXEL*(THREADS_PER_CTA/32)*ELEMENTS_PER_LDG*sizeof(float);
+    int bwd_smem_bytes = SMEM_SIZE_BWD + bwd_reduction_bytes;
+    int occupancy = MaxSharedMemoryPerMultiprocessor(device_id) / bwd_smem_bytes;
+    return std::min(max_cta_per_sm, occupancy);
+  }
+};
+
+const std::vector<size_t> NhwcBatchNormAddRelu::numWorkspaceBytes() const {
+  assert(c_ > 0);
+
+  // choose the max memory required between fwd/bwd passes
+  int grid_x_fwd = div_up(m_, PIXELS_PER_CTA_FWD);
+  int grid_x_bwd = div_up(m_, PIXELS_PER_CTA_BWD);
+  int grid_x = max(grid_x_fwd, grid_x_bwd);
+  int grid_y = div_up(c_, C_ELEMENTS_PER_CTA);
+
+  const size_t num_mean_bytes     = c_ * sizeof(float);
+  const size_t num_variance_bytes = num_mean_bytes;
+
+  int elems_per_group = ((m_ + 31) & ~31) * 2;
+  int group_count = div_up(c_, C_ELEMENTS_PER_CTA);
+  const size_t bitmask_bytes = elems_per_group * group_count * sizeof(unsigned int);
+
+  const size_t size_sums          = grid_y*grid_x*THREADS_PER_PIXEL*\
+      ELEMENTS_PER_LDG*2*sizeof(float);
+  const size_t size_counts        = grid_y*grid_x*sizeof(int);
+
+  return {num_mean_bytes, num_variance_bytes, bitmask_bytes,
+          size_retired_ctas(grid_y), size_sums, size_counts};
+}
+
+void NhwcBatchNormAddRelu::setWorkspacePointers(
+      const std::vector<void*>& workspace,
+      const std::vector<size_t>& num_workspace_bytes) {
+  assert(workspace.size() == 6);
+  assert(num_workspace_bytes.size() == 6);
+
+  minibatch_mean_     = static_cast<float*>(workspace[0]);
+  minibatch_variance_ = static_cast<float*>(workspace[1]);
+  relu_bitmask_       = static_cast<unsigned int*>(workspace[2]);
+  retired_ctas_       = static_cast<int*>(workspace[3]);
+  partial_sums_       = static_cast<float*>(workspace[4]);
+  partial_counts_     = static_cast<int*>(workspace[5]);
+}
+
+void NhwcBatchNormAddRelu::_setFwdParams(NhwcBatchNormFwdParams *params) const {
+  params->gmem_src          = static_cast<uint16_t*>(X_);
+  params->gmem_dst          = static_cast<uint16_t*>(Y_);
+  params->gmem_src1         = static_cast<uint16_t*>(addend_);
+  params->gmem_bias         = bias_;
+  params->gmem_scale        = scale_;
+  params->gmem_running_mean = population_mean_;
+  params->gmem_running_var  = population_variance_;
+  params->gmem_saved_mean   = minibatch_mean_;
+  params->gmem_saved_var    = minibatch_variance_;
+  params->gmem_relu_bitmask = relu_bitmask_;
+  params->nhw               = m_;
+  params->c                 = c_;
+  params->svar_inv_count    = svar_inv_count_;
+  params->rvar_inv_count    = rvar_inv_count_;
+  params->gmem_sums         = partial_sums_;
+  params->gmem_counts       = partial_counts_;
+  params->gmem_retired_ctas = retired_ctas_;
+  params->var_eps           = eps_;
+  params->outer_loops       = 0;
+  params->exp_avg_factor    = static_cast<float>(exp_avg_factor_);
+  params->c_blks            = div_up(c_, C_ELEMENTS_PER_CTA);
+}
+
+void NhwcBatchNormAddRelu::_setFwdInferenceParams(NhwcBatchNormFwdInferenceParams
+                                                        *params) const {
+  params->gmem_src   = static_cast<uint16_t*>(X_);
+  params->gmem_dst   = static_cast<uint16_t*>(Y_);
+  params->gmem_src1  = static_cast<uint16_t*>(addend_);
+  params->gmem_bias  = bias_;
+  params->gmem_scale = scale_;
+  params->gmem_mean  = population_mean_;
+  params->gmem_var   = population_variance_;
+  params->nhw        = m_;
+  params->c          = c_;
+  params->var_eps    = eps_;
+}
+
+void NhwcBatchNormAddRelu::_setBwdParams(NhwcBatchNormBwdParams *params) const {
+  params->gmem_src          = static_cast<uint16_t*>(X_);
+  params->gmem_dy           = static_cast<uint16_t*>(dY_);
+  params->gmem_dst          = static_cast<uint16_t*>(dX_);
+  params->gmem_dst1         = static_cast<uint16_t*>(dAddend_);
+  params->gmem_relu_bitmask = relu_bitmask_;
+  params->gmem_dscale       = dscale_;
+  params->gmem_dbias        = dbias_;
+  params->gmem_scale        = scale_;
+  params->gmem_bias         = bias_;
+  params->gmem_saved_mean   = minibatch_mean_;
+  params->gmem_saved_var    = minibatch_variance_;
+  params->nhw               = m_;
+  params->c                 = c_;
+  params->svar_inv_count    = svar_inv_count_;
+  params->gmem_sums         = partial_sums_;
+  params->gmem_retired_ctas = retired_ctas_;
+  params->outer_loops       = 0;
+  params->c_blks            = div_up(c_, C_ELEMENTS_PER_CTA);
+}
+
+void NhwcBatchNormAddRelu::fwdInference(cudaStream_t stream) {
+  bool ptrs_are_set =
+      X_tensor_desc_ != nullptr
+      && Y_tensor_desc_ != nullptr
+      && scale_ != nullptr
+      && bias_ != nullptr
+      //      && minibatch_mean_ != nullptr
+      //      && minibatch_variance_ != nullptr
+      && population_mean_ != nullptr
+      && population_variance_ != nullptr
+      && X_ != nullptr
+      //      && dX_ != nullptr
+      && Y_ != nullptr
+      && addend_ != nullptr
+      //      && dY_ != nullptr
+      //      && dscale_ != nullptr
+      //      && dbias_ != nullptr
+      && partial_sums_   != nullptr
+      && partial_counts_ != nullptr;
+
+  if (!ptrs_are_set)
+    die();
+
+  dim3 grid_dim;
+  grid_dim.x = div_up(m_, PIXELS_PER_CTA_FWD_INFERENCE);
+  grid_dim.y = div_up(c_, C_ELEMENTS_PER_CTA);
+
+  // @todo: maybe just move this inside initialize routine?
+  NhwcBatchNormFwdInferenceParams params;
+  _setFwdInferenceParams(&params);
+
+  nhwc_batch_norm_fwd_inference
+    <StorageType, THREADS_PER_CTA, THREADS_PER_PIXEL, ELEMENTS_PER_LDG, false, true>
+  <<<grid_dim, THREADS_PER_CTA, 0, stream>>>(params);
+  checkCudaStatus(name_ + " fwd_inference-relu kernel");
+}
+
+dim3 NhwcBatchNormAddRelu::calc_fwd_grid(int *loop, const int grid_dim_x) {
+  dim3 grid_dim;
+  grid_dim.x = div_up(m_, PIXELS_PER_CTA_FWD);
+  int c_blks = div_up(c_, C_ELEMENTS_PER_CTA);
+  unsigned int max_grid_x = grid_dim_x;
+  if (grid_dim.x <= max_grid_x) {
+    *loop = 1;
+    if (max_grid_x / grid_dim.x > 1) {
+      grid_dim.y = std::min(c_blks, static_cast<int>(max_grid_x / grid_dim.x));
+      assert(grid_dim.y<MAX_GBN_BLOCK_Y); //FIXME: turn into a loop
+    } else {
+      grid_dim.y = 1;
+    }
+  } else {
+    grid_dim.x = max_grid_x;
+    grid_dim.y = 1;
+    int nhw_in_regs = m_ - PIXELS_PER_THREAD_IN_SMEM_FWD*PIXELS_PER_LDG*grid_dim.x;
+    int pixels_per_iteration = PIXELS_PER_THREAD_IN_REGISTERS_FWD*PIXELS_PER_LDG*grid_dim.x;
+    *loop = div_up(nhw_in_regs, pixels_per_iteration);
+  }
+  return grid_dim;
+}
+
+dim3 NhwcBatchNormAddRelu::calc_bwd_grid(int *loop, const int grid_dim_x) {
+  dim3 grid_dim;
+  grid_dim.x = div_up(m_, PIXELS_PER_CTA_BWD);
+  int c_blks = div_up(c_, C_ELEMENTS_PER_CTA);
+  unsigned int max_grid_x = grid_dim_x;
+  if (grid_dim.x <= max_grid_x) {
+    *loop = 1;
+    if (max_grid_x / grid_dim.x > 1) {
+      grid_dim.y = std::min(c_blks, static_cast<int>(max_grid_x / grid_dim.x));
+      assert(grid_dim.y<MAX_GBN_BLOCK_Y); //FIXME: turn into a loop
+    } else {
+      grid_dim.y = 1;
+    }
+  } else {
+    grid_dim.x = max_grid_x;
+    grid_dim.y = 1;
+    int nhw_in_regs = m_ - PIXELS_PER_THREAD_IN_SMEM_BWD*PIXELS_PER_LDG*grid_dim.x;
+    int pixels_per_iteration = PIXELS_PER_THREAD_IN_REGISTERS_BWD*PIXELS_PER_LDG*grid_dim.x;
+    *loop = div_up(nhw_in_regs, pixels_per_iteration);
+  }
+  return grid_dim;
+}
+
+void NhwcBatchNormAddRelu::fwd(cudaStream_t stream, void* my_data, void* pair_data, void* pair_data2, void* pair_data3,
+                               const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop) {
+  bool ptrs_are_set =
+      X_tensor_desc_ != nullptr
+      && Y_tensor_desc_ != nullptr
+      && scale_ != nullptr
+      && bias_ != nullptr
+      && minibatch_mean_ != nullptr
+      && minibatch_variance_ != nullptr
+      && relu_bitmask_ != nullptr
+      && population_mean_ != nullptr
+      && population_variance_ != nullptr
+      && X_ != nullptr
+      //      && dX_ != nullptr
+      && Y_ != nullptr
+      && addend_ != nullptr
+      //      && dY_ != nullptr
+      //      && dscale_ != nullptr
+      //      && dbias_ != nullptr
+      && partial_sums_   != nullptr
+      && partial_counts_ != nullptr
+      && retired_ctas_   != nullptr;
+
+  if (!ptrs_are_set)
+    die();
+
+  // reset of retired_cta_count no longer needed
+
+  NhwcBatchNormFwdParams params;
+  _setFwdParams(&params);
+
+  params.my_data = my_data;
+  params.pair_datas[0] = pair_data;
+  params.pair_datas[1] = pair_data2;
+  params.pair_datas[2] = pair_data3;
+  params.magic = magic;
+  params.sync_iters = (bn_group==8)?3:(bn_group >> 1);
+
+  dim3 grid_dim = calc_fwd_grid(&params.outer_loops, grid_dim_x);
+  _fwdKernelLauncher(stream, params, grid_dim, params.outer_loops, occupancy, coop);
+}
+
+void NhwcBatchNormAddRelu::dgrad(cudaStream_t stream, void* my_data, void* pair_data, void* pair_data2, void* pair_data3,
+                                 const int bn_group, const int magic, const int occupancy, const int grid_dim_x, const bool coop) {
+  bool ptrs_are_set =
+      X_tensor_desc_ != nullptr
+      && Y_tensor_desc_ != nullptr
+      && scale_ != nullptr
+      && bias_ != nullptr
+      && minibatch_mean_ != nullptr
+      && minibatch_variance_ != nullptr
+      && relu_bitmask_ != nullptr
+      //      && population_mean_ != nullptr
+      //      && population_variance_ != nullptr
+      && X_ != nullptr
+      && dX_ != nullptr
+      //      && Y_ != nullptr
+      && dY_ != nullptr
+      && dAddend_ != nullptr
+      && dscale_ != nullptr
+      && dbias_ != nullptr
+      && retired_ctas_   != nullptr;
+
+  if (!ptrs_are_set)
+    die();
+
+  // reset of retired_cta_count no longer needed
+
+  NhwcBatchNormBwdParams params;
+  _setBwdParams(&params);
+
+  params.my_data = my_data;
+  params.pair_datas[0] = pair_data;
+  params.pair_datas[1] = pair_data2;
+  params.pair_datas[2] = pair_data3;
+  params.magic = magic;
+  params.sync_iters = (bn_group==8)?3:(bn_group >> 1);
+  params.wgrad_coeff = 1.0 / bn_group;
+
+  dim3 grid_dim = calc_bwd_grid(&params.outer_loops, grid_dim_x);
+  _bwdKernelLauncher(stream, params, grid_dim, params.outer_loops, occupancy, coop);
+}
+
+#endif  // MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_ADD_RELU_H_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/cuda_utils.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/cuda_utils.h
new file mode 100644
index 0000000000..9f003840c0
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/cuda_utils.h
@@ -0,0 +1,20 @@
+#include <ATen/cuda/CUDAContext.h>
+#ifndef CUDA_UTILS_H
+#define CUDA_UTILS_H
+
+namespace at {
+namespace cuda {
+
+namespace utils {
+
+static inline int MaxSharedMemoryPerMultiprocessor(int device_id) {
+    return getDeviceProperties(device_id)->sharedMemPerMultiprocessor;
+}
+
+
+}
+}
+}
+
+
+#endif
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/interface.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/interface.cpp
new file mode 100644
index 0000000000..8cea5f9d93
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/interface.cpp
@@ -0,0 +1,175 @@
+#include <pybind11/pybind11.h>
+#include <pybind11/numpy.h>
+#include <pybind11/stl.h>
+
+#include <torch/extension.h>
+#include <ATen/ATen.h>
+#include <ATen/ArrayRef.h>
+#include <ATen/ScalarType.h>
+#include "ATen/Scalar.h"
+#ifndef VERSION_GE_1_1
+#include "ATen/Type.h"
+#endif
+#include "ATen/Tensor.h"
+#include "ATen/Storage.h"
+#include "ATen/Generator.h"
+
+
+namespace py = pybind11;
+
+int64_t get_buffer_size(
+                       const int bn_sync_steps);
+
+void* get_data_ptr(
+                       const at::Tensor& data);
+
+void* get_remote_data_ptr(
+                       const at::Tensor& handle,
+                       const int64_t offset);
+
+void close_remote_data(
+                       const at::Tensor& handle);
+
+at::Tensor nhwc_bn_fwd_train(
+                       const at::Tensor& x,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& minibatch_mean,
+                       const at::Tensor& minibatch_inv_var,
+                       const at::Tensor& ret_cta,
+                       const float momentum,
+                       const float epsilon,
+                       const bool fuse_relu,
+                       void* my_data,
+                       void* pair_data,
+                       void* pair_data2,
+                       void* pair_data3,
+                       const int bn_group,
+                       const at::Tensor& magic_tensor,
+                       const int occupancy,
+                       const int grid_dim_x,
+                       const bool coop);
+
+at::Tensor nhwc_bn_fwd_eval(
+                       const at::Tensor& x,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& ret_cta,
+                       const int bn_group,
+                       const float momentum,
+                       const float epsilon,
+                       const bool fuse_relu);
+
+std::vector<at::Tensor> nhwc_bn_bwd(
+                       const at::Tensor& x,
+                       const at::Tensor& dy,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& minibatch_mean,
+                       const at::Tensor& minibatch_inv_var,
+                       const at::Tensor& ret_cta,
+                       const float momentum,
+                       const float epsilon,
+                       const bool fuse_relu,
+                       void* my_data,
+                       void* pair_data,
+                       void* pair_data2,
+                       void* pair_data3,
+                       const int bn_group,
+                       const at::Tensor& magic_tensor,
+                       const int occupancy,
+                       const int grid_dim_x,
+                       const bool coop);
+
+at::Tensor nhwc_bn_addrelu_fwd_train(
+                       const at::Tensor& x,
+                       const at::Tensor& z,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& minibatch_mean,
+                       const at::Tensor& minibatch_inv_var,
+                       const at::Tensor& bitmask,
+                       const at::Tensor& ret_cta,
+                       const float momentum,
+                       const float epsilon,
+                       void* my_data,
+                       void* pair_data,
+                       void* pair_data2,
+                       void* pair_data3,
+                       const int bn_group,
+                       const at::Tensor& magic_tensor,
+                       const int occupancy,
+                       const int grid_dim_x,
+                       const bool coop);
+
+at::Tensor nhwc_bn_addrelu_fwd_eval(
+                       const at::Tensor& x,
+                       const at::Tensor& z,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& ret_cta,
+                       const int bn_group,
+                       const float momentum,
+                       const float epsilon);
+
+std::vector<at::Tensor> nhwc_bn_addrelu_bwd(
+                       const at::Tensor& x,
+                       const at::Tensor& dy,
+                       const at::Tensor& scale,
+                       const at::Tensor& bias,
+                       const at::Tensor& running_mean,
+                       const at::Tensor& running_inv_var,
+                       const at::Tensor& minibatch_mean,
+                       const at::Tensor& minibatch_inv_var,
+                       const at::Tensor& bitmask,
+                       const at::Tensor& ret_cta,
+                       const float momentum,
+                       const float epsilon,
+                       void* my_data,
+                       void* pair_data,
+                       void* pair_data2,
+                       void* pair_data3,
+                       const int bn_group,
+                       const at::Tensor& magic_tensor,
+                       const int occupancy,
+                       const int grid_dim_x,
+                       const bool coop);
+
+int nhwc_bn_fwd_occupancy();
+int nhwc_bn_bwd_occupancy();
+
+int nhwc_bn_addrelu_fwd_occupancy();
+int nhwc_bn_addrelu_bwd_occupancy();
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+
+  m.def("get_buffer_size", &get_buffer_size, "get_buffer_size");
+  m.def("get_data_ptr", &get_data_ptr, "get_data_ptr");
+  m.def("get_remote_data_ptr", &get_remote_data_ptr, "get_remote_data_ptr");
+  m.def("close_remote_data", &close_remote_data, "close_remote_data");
+
+  m.def("bn_fwd_nhwc", &nhwc_bn_fwd_train, "bn_fwd_nhwc");
+  m.def("bn_fwd_eval_nhwc", &nhwc_bn_fwd_eval, "bn_fwd_eval_nhwc");
+  m.def("bn_bwd_nhwc", &nhwc_bn_bwd, "bn_bwd_nhwc");
+
+  m.def("bn_fwd_nhwc_occupancy", &nhwc_bn_fwd_occupancy, "bn_fwd_nhwc_occupancy");
+  m.def("bn_bwd_nhwc_occupancy", &nhwc_bn_bwd_occupancy, "bn_bwd_nhwc_occupancy");
+
+  m.def("bn_addrelu_fwd_nhwc", &nhwc_bn_addrelu_fwd_train, "bn_addrelu_fwd_nhwc");
+  m.def("bn_addrelu_fwd_eval_nhwc", &nhwc_bn_addrelu_fwd_eval, "bn_addrelu_fwd_eval_nhwc");
+  m.def("bn_addrelu_bwd_nhwc", &nhwc_bn_addrelu_bwd, "bn_addrelu_bwd_nhwc");
+
+  m.def("bn_addrelu_fwd_nhwc_occupancy", &nhwc_bn_addrelu_fwd_occupancy, "bn_addrelu_fwd_nhwc_occupancy");
+  m.def("bn_addrelu_bwd_nhwc_occupancy", &nhwc_bn_addrelu_bwd_occupancy, "bn_addrelu_bwd_nhwc_occupancy");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/ipc.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/ipc.cu
new file mode 100644
index 0000000000..6b152a0d0d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/ipc.cu
@@ -0,0 +1,129 @@
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <cuda.h>
+
+#include "compat.h"
+
+
+#define cudaCheckErrors(msg) \
+    do { \
+        cudaError_t __err = cudaGetLastError(); \
+        if (__err != cudaSuccess) { \
+            fprintf(stderr, "Fatal error: %s (%s at %s:%d)\n", \
+                msg, cudaGetErrorString(__err), \
+                __FILE__, __LINE__); \
+            fprintf(stderr, "*** FAILED - ABORTING\n"); \
+            exit(1); \
+        } \
+    } while (0)
+
+template<>
+struct std::hash<cudaIpcMemHandle_t> {
+  size_t operator() (const cudaIpcMemHandle_t& handle) const {
+    size_t hash = 0;
+    uint8_t* ptr = (uint8_t*)&handle;
+    assert(sizeof(uint8_t) == 1);
+    for (int i=0; i<sizeof(cudaIpcMemHandle_t); i++) {
+      hash += *ptr;
+      ptr++;
+    }
+    return hash;
+  }
+};
+
+template<>
+struct std::equal_to<cudaIpcMemHandle_t> {
+  bool operator() (const cudaIpcMemHandle_t &lhs,
+                             const cudaIpcMemHandle_t &rhs) const {
+    return (std::memcmp((void*) &lhs,
+                        (void*) &rhs,
+                        sizeof(cudaIpcMemHandle_t)) == 0);
+  }
+};
+
+namespace {
+
+namespace gpuipc {
+//from: src/operator/nn/cudnn/nhwc_batch_norm_kernel.h
+// The number of threads per pixel.
+const int THREADS_PER_PIXEL = 16;
+// The number of elements per ldg.
+const int ELEMENTS_PER_LDG = 4;
+// The number of reducing ops, each uses its own space : mean, var, dscale, dbias
+const int REDUCE_OPS = 4;
+// Maximum block.y supported - limited due to buffer allocation
+const int MAX_BLOCK_Y = 256;
+const int MAX_OFFSET = REDUCE_OPS*MAX_BLOCK_Y;
+const int BYTES_PER_ELEM = 4;
+// Buffer size per sync step
+const int SINGLE_SYNC_BUFFER_BYTES = MAX_OFFSET*THREADS_PER_PIXEL*2*ELEMENTS_PER_LDG*BYTES_PER_ELEM;
+};
+
+class IpcMemHandleRegistry {
+public:
+  void* getPtr(const cudaIpcMemHandle_t& handle, int64_t offset) {
+    if (registry_.count(handle) == 0) {
+      registry_.insert(std::make_pair(handle, RegistryEntry()));
+      registry_[handle].dev_ptr = ipcOpenMem(handle);
+    }
+    registry_[handle].ref_count++;
+    return (((uint8_t*)registry_[handle].dev_ptr) + offset);
+  }
+
+  void releasePtr(const cudaIpcMemHandle_t& handle) {
+    if (registry_.count(handle) == 0) {
+    }
+    if (--registry_[handle].ref_count == 0) {
+      ipcCloseMem(registry_[handle].dev_ptr);
+      registry_.erase(handle);
+    }
+  }
+
+  struct RegistryEntry {
+    void* dev_ptr;
+    int   ref_count;
+    RegistryEntry() : dev_ptr(NULL) , ref_count(0) {}
+  };
+
+protected:
+  std::unordered_map<cudaIpcMemHandle_t, RegistryEntry> registry_;
+
+  void* ipcOpenMem(const cudaIpcMemHandle_t& handle) {
+    void *data;
+    cudaIpcOpenMemHandle(&data, handle, cudaIpcMemLazyEnablePeerAccess);
+    cudaCheckErrors("ipc init");
+    return data;
+  }
+
+  void ipcCloseMem(void* dev_ptr) {
+    cudaIpcCloseMemHandle(dev_ptr);
+    cudaCheckErrors("ipc close");
+  }
+
+};
+
+}
+
+static IpcMemHandleRegistry ipc_mem_registry;
+
+int64_t get_buffer_size(const int bn_sync_steps) {
+  return bn_sync_steps * gpuipc::SINGLE_SYNC_BUFFER_BYTES;
+}
+
+void* get_remote_data_ptr(const at::Tensor& handle, const int64_t offset) {
+  cudaIpcMemHandle_t my_handle;
+  memcpy((unsigned char *)(&my_handle), handle.DATA_PTR<uint8_t>(), sizeof(my_handle));
+  return ipc_mem_registry.getPtr(my_handle, offset);
+}
+
+void close_remote_data(const at::Tensor& handle) {
+    cudaIpcMemHandle_t my_handle;
+    memcpy((unsigned char *)(&my_handle), handle.DATA_PTR<uint8_t>(), sizeof(my_handle));
+  ipc_mem_registry.releasePtr(my_handle);
+}
+
+void* get_data_ptr(
+                   const at::Tensor& data) {
+  return data.DATA_PTR<uint8_t>();
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/nhwc_batch_norm_kernel.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/nhwc_batch_norm_kernel.h
new file mode 100644
index 0000000000..8430f30994
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/groupbn/nhwc_batch_norm_kernel.h
@@ -0,0 +1,2685 @@
+/*
+ * Licensed to the Apache Software Foundation (ASF) under one
+ * or more contributor license agreements.  See the NOTICE file
+ * distributed with this work for additional information
+ * regarding copyright ownership.  The ASF licenses this file
+ * to you 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.
+ */
+
+/*!
+ * Copyright (c) 2018 by Contributors
+ * \file nhwc_batch_norm_kernel.h
+ * \brief CUDA NHWC Batch Normalization code
+ * \author Shankara Rao Thejaswi Nanditale, Dick Carter, Maxim Milakov, Evgeni Krimer
+*/
+#ifndef MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_KERNEL_H_
+#define MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_KERNEL_H_
+
+#include <stdint.h>
+#include <algorithm>
+
+#define DEVICE_FUNCTION static inline __device__
+
+// CTA margin used by cooperative launch. Can be overridden by env var NHWC_BATCHNORM_LAUNCH_MARGIN.
+#define NHWC_BATCHNORM_LAUNCH_MARGIN_MIN     3
+#define NHWC_BATCHNORM_LAUNCH_MARGIN_DEFAULT NHWC_BATCHNORM_LAUNCH_MARGIN_MIN
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< typename T, int ELEMENTS_PER_LDG >
+struct PackedStorage {
+    enum { PACKED_ELEMENTS_PER_LDG = ELEMENTS_PER_LDG };
+    typedef T Type;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int ELEMENTS_PER_LDG >
+struct PackedStorage<uint16_t, ELEMENTS_PER_LDG> {
+    enum { PACKED_ELEMENTS_PER_LDG = ELEMENTS_PER_LDG/2 };
+    typedef int Type;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+DEVICE_FUNCTION void from_float(int (&dst)[N], const float (&src)[2*N]) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        uint16_t lo, hi;
+        asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(lo) : "f"(src[2*i+0]));
+        asm volatile("cvt.rn.f16.f32 %0, %1;" : "=h"(hi) : "f"(src[2*i+1]));
+        asm volatile("mov.b32 %0, {%1, %2};"  : "=r"(dst[i]) : "h"(lo), "h"(hi));
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+DEVICE_FUNCTION void from_float(float (&dst)[N], const float (&src)[N]) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        dst[i] = src[i];
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+DEVICE_FUNCTION void to_float(float (&dst)[2*N], int (&src)[N]) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        uint16_t lo, hi;
+        asm volatile("mov.b32 {%0, %1}, %2;" : "=h"(lo), "=h"(hi) : "r"(src[i]));
+        asm volatile("cvt.f32.f16 %0, %1;"   : "=f"(dst[2*i+0])   : "h"(lo));
+        asm volatile("cvt.f32.f16 %0, %1;"   : "=f"(dst[2*i+1])   : "h"(hi));
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+DEVICE_FUNCTION void to_float(float (&dst)[N], float (&src)[N]) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        dst[i] = src[i];
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void ldg(int (&dst)[1], const uint16_t *gmem) {
+    dst[0] = __ldg((const int*) gmem);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void ldg_stream(int (&dst)[1], const uint16_t *gmem) {
+    unsigned int tmp;
+    asm volatile ("ld.global.cs.nc.s32 %0, [%1];"  : "=r"(tmp) : "l" ((const uint *)gmem));
+    dst[0] = tmp;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void ldg(int (&dst)[2], const uint16_t *gmem) {
+    int2 tmp = __ldg((const int2*) gmem);
+    dst[0] = tmp.x;
+    dst[1] = tmp.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void ldg_stream(int (&dst)[2], const uint16_t *gmem) {
+    int2 tmp;
+    asm volatile ("ld.global.cs.nc.v2.s32 {%0,%1}, [%2];"
+        : "=r"(tmp.x), "=r"(tmp.y) : "l"((const int2 *)gmem));
+    dst[0] = tmp.x;
+    dst[1] = tmp.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+DEVICE_FUNCTION void ldg(float (&dst)[N], const uint16_t *gmem) {
+    int tmp[N/2];
+    ldg(tmp, gmem);
+    to_float(dst, tmp);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+DEVICE_FUNCTION void ldg_stream(float (&dst)[N], const uint16_t *gmem) {
+    int tmp[N/2];
+    ldg_stream(tmp, gmem);
+    to_float(dst, tmp);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void stg(uint16_t *gmem, int (&src)[1]) {
+    reinterpret_cast<int*>(gmem)[0] = src[0];
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void stg_stream(uint16_t *gmem, int (&src)[1]) {
+    unsigned int tmp = src[0];
+    asm volatile ("st.global.cs.s32 [%0], %1;"
+        :: "l"((uint *)gmem) , "r"(tmp));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void stg(uint16_t *gmem, int (&src)[2]) {
+    reinterpret_cast<int2*>(gmem)[0] = make_int2(src[0], src[1]);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void stg_stream(uint16_t *gmem, int (&src)[2]) {
+    asm volatile ("st.global.cs.v2.s32 [%0], {%1,%2};"
+        :: "l"((uint *)gmem) , "r"(src[0]), "r"( src[1]));
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+DEVICE_FUNCTION void stg(uint16_t *gmem, float (&src)[N]) {
+    int tmp[N/2];
+    from_float(tmp, src);
+    stg(gmem, tmp);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+DEVICE_FUNCTION void stg_stream(uint16_t *gmem, float (&src)[N]) {
+    int tmp[N/2];
+    from_float(tmp, src);
+    stg_stream(gmem, tmp);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void read_from_gmem(float (&dst)[2], const float *gmem, int idx) {
+    float2 tmp = __ldg(reinterpret_cast<const float2*>(&gmem[2*idx]));
+    dst[0] = tmp.x;
+    dst[1] = tmp.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void read_from_gmem(float (&dst)[4], const float *gmem, int idx) {
+    float4 tmp = __ldg(reinterpret_cast<const float4*>(&gmem[4*idx]));
+    dst[0] = tmp.x;
+    dst[1] = tmp.y;
+    dst[2] = tmp.z;
+    dst[3] = tmp.w;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void read_from_smem(float (&x)[2], const float *smem, int idx) {
+    float2 tmp = *(const float2*) &smem[2*idx];
+    x[0] = tmp.x;
+    x[1] = tmp.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void read_from_smem(int (&x)[1], const int *smem, int idx) {
+    x[0] = smem[idx];
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void read_from_smem(float (&x)[4], const float *smem, int idx) {
+    float4 tmp = *(const float4*) &smem[4*idx];
+    x[0] = tmp.x;
+    x[1] = tmp.y;
+    x[2] = tmp.z;
+    x[3] = tmp.w;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void read_from_smem(int (&x)[2], const int *smem, int idx) {
+    int2 tmp = *(const int2*) &smem[2*idx];
+    x[0] = tmp.x;
+    x[1] = tmp.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void write_to_gmem(float *gmem, int idx, const float (&src)[2]) {
+    reinterpret_cast<float2*>(&gmem[2*idx])[0] = make_float2(src[0], src[1]);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void write_to_gmem(float *gmem, int idx, const float (&src)[4]) {
+    reinterpret_cast<float4*>(&gmem[4*idx])[0] = make_float4(src[0], src[1], src[2], src[3]);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void scaled_write_to_gmem(float *gmem, int idx, const float (&src)[4], const float coeff) {
+    reinterpret_cast<float4*>(&gmem[4*idx])[0] = make_float4(src[0]*coeff, src[1]*coeff, src[2]*coeff, src[3]*coeff);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void write_to_smem(float *smem, int idx, const float (&x)[2]) {
+    reinterpret_cast<float2*>(&smem[2*idx])[0] = make_float2(x[0], x[1]);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void write_to_smem(int *smem, int idx, const int (&x)[1]) {
+    smem[idx] = x[0];
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void write_to_smem(float *smem, int idx, const float (&x)[4]) {
+    reinterpret_cast<float4*>(&smem[4*idx])[0] = make_float4(x[0], x[1], x[2], x[3]);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+DEVICE_FUNCTION void write_to_smem(int *smem, int idx, const int (&x)[2]) {
+    reinterpret_cast<int2*>(&smem[2*idx])[0] = make_int2(x[0], x[1]);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+DEVICE_FUNCTION void zero_array(int (&dst)[N]) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        dst[i] = 0;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int N >
+DEVICE_FUNCTION void zero_array(float (&dst)[N]) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        dst[i] = 0.f;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void add(float (&x)[N], const float (&y)[N]) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        x[i] += y[i];
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void multiply(float (&x)[N], const float (&y)[N]) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        x[i] *= y[i];
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void scale_(float (&x)[N], float scalar) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        x[i] *= scalar;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void normalize(float (&x)[N], const float (&bias)[N],
+                               const float (&scale)[N], const float (&m1)[N]) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        x[i] = bias[i] + scale[i] * (x[i] - m1[i]);
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename Storage>
+DEVICE_FUNCTION Storage relu(Storage in) {
+    Storage zero = (Storage)0.f;
+    return (in < zero)? zero : in;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void relu_activation(float (&x)[N]) {
+    #pragma unroll
+    for (int i = 0; i < N; ++i) {
+        x[i] = relu(x[i]);
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+template< int THREADS_PER_CTA >
+DEVICE_FUNCTION void parallel_sums_16x2(float *smem, float (&x)[4], int nhw,
+                                        void* params_my_data, void** params_pair_datas, int off,
+                                        const int magic,
+                                        const int sync_iters) {
+    // The size of a warp.
+    const int THREADS_PER_WARP = 32;
+    // The number of warps in a CTA.
+    const int WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;
+    // The number of threads per pixel.
+    const int THREADS_PER_PIXEL = 16;
+    // The number of elements per ldg.
+    const int ELEMENTS_PER_LDG = 4;
+    // The number of reducing ops, each uses its own space : mean, var, dscale, dbias
+    const int REDUCE_OPS = 4;
+    // Maximum block.y supported - limited due to buffer allocation
+    const int MAX_BLOCK_Y = 256;
+    const int MAX_OFFSET = REDUCE_OPS*MAX_BLOCK_Y;
+    // The warp decomposition.
+    const int warp_id = threadIdx.x / THREADS_PER_WARP;
+    const int lane_id = threadIdx.x % THREADS_PER_WARP;
+    // total size of data per sync iter
+    const int data_total = MAX_OFFSET*THREADS_PER_PIXEL*ELEMENTS_PER_LDG*2;
+
+    #pragma unroll
+    for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+        x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL+lane_id);
+    }
+
+    // The warp leaders, write to SMEM.
+    if (lane_id < THREADS_PER_PIXEL) {
+        write_to_smem(smem, warp_id*THREADS_PER_PIXEL + lane_id, x);
+    }
+
+    // The data is in SMEM. Do the final reduction.
+    __syncthreads();
+
+    // The 1st warp does all the work.
+    // We do the final reduction each half-warp sequentially reduces the final values.
+    if (warp_id == 0) {
+        read_from_smem(x, smem, threadIdx.x);
+
+        #pragma unroll
+        for (int offset = 1;
+             offset < WARPS_PER_CTA/(THREADS_PER_WARP / THREADS_PER_PIXEL); ++offset) {
+            float y[ELEMENTS_PER_LDG];
+            // Read the mean and variance from the other pixel.
+            read_from_smem(y, smem, threadIdx.x + offset*THREADS_PER_WARP);
+            // Compute the updated sum.
+            add(x, y);
+        }
+
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL+lane_id);
+        }
+
+        // Make sure the data was read from SMEM.
+        __syncwarp();
+
+        // Store the final values.
+        if (threadIdx.x < THREADS_PER_PIXEL) {
+        // probably could do it earlier, before sync
+
+        for (int sync_iter=0; sync_iter < sync_iters; ++sync_iter) {
+            //float* params_pair_data = (reinterpret_cast<float**>(params_pair_datas))[sync_iter];
+            void* params_pair_data = params_pair_datas[sync_iter];
+
+            // skip the space consumed by previous sync iterations
+            const int xbuf_offset = sync_iter*data_total;
+            // data starts after flags, but have to skip previous
+            const int data_offset = xbuf_offset
+                                    + off*ELEMENTS_PER_LDG*THREADS_PER_PIXEL*2
+                                    + ELEMENTS_PER_LDG*threadIdx.x*2;
+
+            // after sums for this GPU were computed, let CTA0 broadcast the sum to over GPU
+            if (blockIdx.x == 0) {
+                volatile float * write_data =
+                    &((reinterpret_cast<float*>(params_pair_data))[data_offset]);
+
+                // write the data to memory region to be reflected to other GPU
+                asm volatile ("st.global.wt.v4.b32 [%0], {%1,%2,%3,%4};"
+                    :: "l"(write_data) , "f"(x[0]), "r"(magic), "f"(x[2]), "r"(magic));
+
+                asm volatile ("st.global.wt.v4.b32 [%0], {%1,%2,%3,%4};"
+                    :: "l"(write_data+4) , "f"(x[1]), "r"(magic), "f"(x[3]), "r"(magic));
+            }
+
+            // now each CTA (on each GPU) reads the data written by CTA 0 of the other GPU
+            volatile float * read_data =
+                &((reinterpret_cast<float*>(params_my_data))[data_offset]);
+
+            float other[4];
+            uint32_t other_flag_a, other_flag_b;
+            do {
+                asm volatile ("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];"
+                    : "=f"(other[0]), "=r"(other_flag_a), "=f"(other[2]), "=r"(other_flag_b) : "l"(read_data));
+            } while ((other_flag_a != magic) || (other_flag_b != magic));
+
+            do {
+                asm volatile ("ld.volatile.global.v4.b32 {%0, %1, %2, %3}, [%4];"
+                    : "=f"(other[1]), "=r"(other_flag_a), "=f"(other[3]), "=r"(other_flag_b) : "l"(read_data+4));
+            } while ((other_flag_a != magic) || (other_flag_b != magic));
+
+            add(x, other);
+        }
+        // finally, after syncing up and accounting for partial sums from
+        // other GPUs as required, write the result
+
+
+            write_to_smem(smem, threadIdx.x, x);
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int THREADS_PER_CTA >
+DEVICE_FUNCTION void parallel_sums_8x4(float *smem, float (&x)[4], int nhw) {
+    // The size of a warp.
+    const int THREADS_PER_WARP = 32;
+    // The number of warps in a CTA.
+    const int WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;
+    // The number of threads per pixel.
+    const int THREADS_PER_PIXEL = 8;
+    // The number of elements per ldg.
+    const int ELEMENTS_PER_LDG = 4;
+    // The warp decomposition.
+    const int warp_id = threadIdx.x / THREADS_PER_WARP;
+    const int lane_id = threadIdx.x % THREADS_PER_WARP;
+
+    #pragma unroll
+    for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+        x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL+lane_id);
+        x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL*2+lane_id);
+    }
+
+    // The warp leaders, write to SMEM.
+    if (lane_id < THREADS_PER_PIXEL) {
+        write_to_smem(smem, warp_id*THREADS_PER_PIXEL + lane_id, x);
+    }
+
+    // The data is in SMEM. Do the final reduction.
+    __syncthreads();
+
+    // The 1st warp does all the work.
+    // We do the final reduction each half-warp sequentially reduces the final values.
+    if (warp_id == 0) {
+        read_from_smem(x, smem, threadIdx.x);
+
+        #pragma unroll
+        for (int offset = 1;
+             offset < WARPS_PER_CTA/(THREADS_PER_WARP / THREADS_PER_PIXEL); ++offset) {
+            float y[ELEMENTS_PER_LDG];
+            // Read the mean and variance from the other pixel.
+            read_from_smem(y, smem, threadIdx.x + offset*THREADS_PER_WARP);
+            // Compute the updated sum.
+            add(x, y);
+        }
+
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL+lane_id);
+            x[i] += __shfl_sync(0xffffffffU, x[i], THREADS_PER_PIXEL*2+lane_id);
+        }
+
+        // Make sure the data was read from SMEM.
+        __syncwarp();
+
+        // Store the final values.
+        if (threadIdx.x < THREADS_PER_PIXEL) {
+            write_to_smem(smem, threadIdx.x, x);
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int THREADS_PER_CTA, int THREADS_PER_PIXEL, int ELEMENTS_PER_LDG >
+DEVICE_FUNCTION void parallel_sums(float *smem, float (&x)[ELEMENTS_PER_LDG], int nhw) {
+    // The size of a warp.
+    const int THREADS_PER_WARP = 32;
+    // The number of warps in a CTA.
+    const int WARPS_PER_CTA = THREADS_PER_CTA / THREADS_PER_WARP;
+    // The number of pixels computed by a single warp.
+    const int PIXELS_PER_WARP = THREADS_PER_WARP / THREADS_PER_PIXEL;
+
+    // The position in the warp.
+    const int nhw_in_warp = nhw % PIXELS_PER_WARP;
+    // The C in the warp.
+    const int c_in_warp = threadIdx.x % THREADS_PER_PIXEL;
+
+    // Store the values to shared memory.
+    write_to_smem(smem, threadIdx.x, x);
+
+    // Compute the parallel sums.
+    for (int offset = PIXELS_PER_WARP/2; offset > 0; offset /= 2) {
+        // NOP.
+        __syncwarp();
+
+        // Read the running sum from the other thread.
+        float y[ELEMENTS_PER_LDG];
+        if (nhw_in_warp < offset) {
+            read_from_smem(y, smem, threadIdx.x + offset*THREADS_PER_PIXEL);
+        }
+
+        // Compute the updated sum.
+        add(x, y);
+
+        // NOP.
+        __syncwarp();
+
+        // Update the sum in SMEM.
+        if (offset > 1 && nhw_in_warp < offset) {
+            write_to_smem(smem, threadIdx.x, x);
+        }
+    }
+
+    // The warps are done. Do the final reduction at the CTA level.
+    __syncthreads();
+
+    // The warp leaders, write to SMEM.
+    const int idx = (threadIdx.x/THREADS_PER_WARP)*THREADS_PER_PIXEL + c_in_warp;
+    if (nhw_in_warp == 0) {
+        write_to_smem(smem, idx, x);
+    }
+
+    // The data is in SMEM. Do the final reduction.
+    __syncthreads();
+
+    // Read the 1st element to prepare the work.
+    if (nhw < WARPS_PER_CTA/2) {
+        read_from_smem(x, smem, threadIdx.x);
+    }
+
+    // We have the running mean and running m2. Let's build the mean/var of the CTA.
+    for (int offset = WARPS_PER_CTA/2; offset > 0; offset /= 2) {
+        // NOP.
+        __syncwarp();
+
+        // Read the mean and variance from the other pixel.
+        float y[ELEMENTS_PER_LDG];
+        if (nhw < offset) {
+            read_from_smem(y, smem, threadIdx.x + offset*THREADS_PER_PIXEL);
+        }
+
+        // Compute the updated sum.
+        add(x, y);
+
+        // NOP.
+        __syncwarp();
+
+        // Store the mean/var for the different pixels.
+        if (nhw < offset) {
+            write_to_smem(smem, threadIdx.x, x);
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template< int THREADS_PER_PIXEL, int ELEMENTS_PER_LDG >
+struct ParallelSums {
+    template< int THREADS_PER_CTA >
+    DEVICE_FUNCTION void dispatch(float *smem, float (&x)[ELEMENTS_PER_LDG], int nhw) {
+        parallel_sums<THREADS_PER_CTA, THREADS_PER_PIXEL, ELEMENTS_PER_LDG>(smem, x, nhw);
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<>
+struct ParallelSums<16, 4> {
+    template< int THREADS_PER_CTA >
+    DEVICE_FUNCTION void dispatch(float *smem, float (&x)[4], int nhw) {
+        parallel_sums_16x2<THREADS_PER_CTA>(smem, x, nhw, 0, 0, 0, 0, 0);
+    }
+
+    template< int THREADS_PER_CTA >
+    DEVICE_FUNCTION void dispatchX(float *smem, float (&x)[4], int nhw, void* params_my_data, void** params_pair_datas, int off, const int magic, const unsigned int& sync_iters) {
+        parallel_sums_16x2<THREADS_PER_CTA>(smem, x, nhw, params_my_data, params_pair_datas, off, magic, sync_iters);
+    }
+};
+
+template<>
+struct ParallelSums<8, 4> {
+    template< int THREADS_PER_CTA >
+    DEVICE_FUNCTION void dispatch(float *smem, float (&x)[4], int nhw) {
+        parallel_sums_8x4<THREADS_PER_CTA>(smem, x, nhw);
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+static inline int div_up(int m, int n) {
+    return (m + n - 1) / n;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// It is expected that all threads in the CTA enter this function!
+DEVICE_FUNCTION void inter_block_sync(int* gmem_retired_ctas, int expected_count, bool master) {
+
+    // Register the CTA.
+    if (threadIdx.x == 0) {
+        // Issue the membar.
+        __threadfence();
+        // Notify that the CTA is done.
+        int val_to_add = 1;
+        if (master) {
+            val_to_add = -(expected_count - 1);
+        }
+        atomicAdd(gmem_retired_ctas, val_to_add);
+    }
+
+    // Are all CTAs done?
+    if (threadIdx.x == 0) {
+        int retired_ctas = -1;
+        do {
+            __threadfence();
+            asm volatile ("ld.global.cg.b32 %0, [%1];"
+                : "=r"(retired_ctas) : "l"(gmem_retired_ctas));
+        } while (retired_ctas != 0);
+    }
+    __syncthreads();
+
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct NhwcBatchNormFwdInferenceParams {
+    // The input/output tensors.
+    uint16_t *gmem_src, *gmem_dst, *gmem_src1;
+    // the final mean and variance as calculated during the training process
+    float *gmem_mean, *gmem_var;
+    // The bias/scale.
+    float *gmem_bias, *gmem_scale;
+    // The dimensions.
+    int nhw, c;
+    // epsilon
+    float var_eps;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// No DESIRED_OCCUPANCY launch bounds needed, as this is not launched cooperatively
+template<
+    typename Storage,
+    int THREADS_PER_CTA,
+    int THREADS_PER_PIXEL,
+    int ELEMENTS_PER_LDG,
+    bool USE_RELU,
+    bool USE_ADD_RELU
+>
+__global__ __launch_bounds__(THREADS_PER_CTA)
+    void nhwc_batch_norm_fwd_inference(NhwcBatchNormFwdInferenceParams params) {
+    // The number of pixels loaded in a single LDG.
+    const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;
+    // The number of C elements per CTA.
+    const int C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL*ELEMENTS_PER_LDG;
+
+    // The start position in the NHW dimension where the CTA starts.
+    const int cta_nhw_stride = gridDim.x * PIXELS_PER_LDG;
+    // Compute the NHW coordinate of the thread in the CTA.
+    const int thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;
+    // thread's starting point in NHW
+    const int thread_nhw = thread_in_cta_nhw + blockIdx.x * PIXELS_PER_LDG;
+
+    // The position in the C dimension where the CTA starts.
+    const int cta_c = blockIdx.y * C_ELEMENTS_PER_CTA;
+    // Compute the C coordinate of the thread in the CTA.
+    const int thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;
+    // Compute the C coordinate of the thread.
+    const int thread_c = cta_c + thread_in_cta_c*ELEMENTS_PER_LDG;
+
+    // Is the thread working on a valid C dimension?
+    const int is_valid_c = thread_c < params.c;
+
+    float mean[ELEMENTS_PER_LDG], var[ELEMENTS_PER_LDG];
+    float scale[ELEMENTS_PER_LDG], bias[ELEMENTS_PER_LDG];
+    zero_array(mean);
+    zero_array(var);
+    zero_array(scale);
+    zero_array(bias);
+    if (is_valid_c) {
+        read_from_gmem(var, &params.gmem_var[cta_c], thread_in_cta_c);
+        read_from_gmem(scale, &params.gmem_scale[cta_c], thread_in_cta_c);
+        read_from_gmem(mean, &params.gmem_mean[cta_c], thread_in_cta_c);
+        read_from_gmem(bias, &params.gmem_bias[cta_c], thread_in_cta_c);
+    }
+
+    // Update the scale with the stddev and eps.
+    #pragma unroll
+    for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+        scale[i] *= rsqrtf(var[i] + params.var_eps);
+    }
+
+    // The base pointers for reading/writing
+    uint16_t *const gmem_src = &params.gmem_src[thread_c];
+    uint16_t *const gmem_dst = &params.gmem_dst[thread_c];
+    const uint16_t *gmem_src1 = nullptr;
+    if (USE_ADD_RELU) {
+        gmem_src1 = &params.gmem_src1[thread_c];
+    }
+
+    // apply BN
+    for (int nhw = thread_nhw; nhw < params.nhw; nhw += cta_nhw_stride) {
+        float x_math[ELEMENTS_PER_LDG];
+        zero_array(x_math);
+        if (is_valid_c) {
+            ldg(x_math, &gmem_src[nhw*params.c]);
+        }
+
+        // Normalize and apply activation function
+        normalize(x_math, bias, scale, mean);
+        if (USE_ADD_RELU) {
+            float x1_math[ELEMENTS_PER_LDG];
+            ldg(x1_math, &gmem_src1[nhw*params.c]);
+            add(x_math, x1_math);
+            relu_activation(x_math);
+        } else if (USE_RELU) {
+            relu_activation(x_math);
+        }
+
+        if (is_valid_c) {
+            stg(&gmem_dst[nhw*params.c], x_math);
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct NhwcBatchNormFwdParams {
+    // The input/output tensors.
+    uint16_t *gmem_src, *gmem_dst, *gmem_src1;
+    // The bias/scale.
+    float *gmem_bias, *gmem_scale;
+    // running mean/var (refer BN API from cudnn doc)
+    float *gmem_running_mean, *gmem_running_var;
+    // saved mean/var (refer BN API from cudnn doc)
+    float *gmem_saved_mean, *gmem_saved_var;
+    // ReLU bitmask
+    unsigned int *gmem_relu_bitmask;
+    // The dimensions.
+    int nhw, c;
+    // factor to scale sum of squared errors to get saved variance.  Must be 1/nhw.
+    float svar_inv_count;
+    // factor to scale sum of squared errors to get running variance. Should be 1/nhw or 1/(nhw-1).
+    float rvar_inv_count;
+    // The buffer to do the reduction for mean, stddev and count.
+    float *gmem_sums;
+    // The buffer to count items in the different CTAs.
+    int *gmem_counts;
+    // The counters of retired CTAs.
+    int *gmem_retired_ctas;
+    // The epsilon to apply to the computation of the variance.
+    float var_eps;
+    // outer loop count
+    int outer_loops;
+    // exponential average factor
+    float exp_avg_factor;
+    // number of CTAs along .x dimension
+    int c_blks;
+
+    void* my_data;
+    void* pair_datas[4];
+    int magic;
+    int sync_iters;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    typename Storage,
+    int THREADS_PER_CTA,
+    int THREADS_PER_PIXEL,
+    int PIXELS_PER_THREAD_IN_REGISTERS,
+    int PIXELS_PER_THREAD_IN_SMEM,
+    int ELEMENTS_PER_LDG,
+    int USE_ONLINE_APPROACH,
+    int OUTER_LOOPS_,
+    bool USE_RELU,
+    bool USE_ADD_RELU,
+    int DESIRED_OCCUPANCY
+>
+__global__ __launch_bounds__(THREADS_PER_CTA, DESIRED_OCCUPANCY)
+    void nhwc_batch_norm_fwd(NhwcBatchNormFwdParams params) {
+    // The number of pixels loaded in a single LDG.
+    const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;
+    // The number of pixels computed per CTA stored in registers.
+    const int PIXELS_PER_CTA_IN_REGISTERS = PIXELS_PER_THREAD_IN_REGISTERS * PIXELS_PER_LDG;
+    // The number of pixels computed per CTA stored in SMEM.
+    const int PIXELS_PER_CTA_IN_SMEM = PIXELS_PER_THREAD_IN_SMEM*PIXELS_PER_LDG;
+    // The number of C elements per CTA.
+    const int C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL*ELEMENTS_PER_LDG;
+
+    // Shared memory to do CTA-wide parallel sums.
+    __shared__ float smem[THREADS_PER_PIXEL*(THREADS_PER_CTA/32)*ELEMENTS_PER_LDG];
+
+    // Compute the NHW coordinate of the thread in the CTA.
+    const int thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;
+
+    // The adapter for the storage.
+    typedef PackedStorage<Storage, ELEMENTS_PER_LDG> PackedStorage_;
+    // The data type for packed storage in SMEM.
+    typedef typename PackedStorage_::Type PackedStorageType;
+    // The number of elements in the packed storage.
+    const int PACKED_ELEMENTS_PER_LDG = PackedStorage_::PACKED_ELEMENTS_PER_LDG;
+    // Registers to keep the data live for the persistent approach.
+    PackedStorageType x_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];
+
+    // Shared memory buffer to store the extra pixels.
+    extern __shared__ PackedStorageType smem_storage_packed[];
+
+    for (int c_blk_index = blockIdx.y; c_blk_index < params.c_blks; c_blk_index += gridDim.y) {
+        // The position in the NHW dimension where the CTA starts.
+        int cta_nhw_regs = blockIdx.x * PIXELS_PER_CTA_IN_REGISTERS;
+        // The position in the NHW dimension where the CTA starts for the portion in SMEM.
+        int cta_nhw_smem = blockIdx.x * PIXELS_PER_CTA_IN_SMEM;
+
+        // The position in the C dimension where the CTA starts.
+        const int cta_c = c_blk_index * C_ELEMENTS_PER_CTA;
+        // Compute the C coordinate of the thread in the CTA.
+        const int thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;
+        // Compute the C coordinate of the thread.
+        int thread_c = cta_c + thread_in_cta_c*ELEMENTS_PER_LDG;
+
+        // Is the thread working on a valid C dimension?
+        const int is_valid_c = thread_c < params.c;
+
+        // Clamp thread_c so that we load from valid locations even if we don't use the value
+        if (!is_valid_c)
+            thread_c = params.c - 4;
+
+        // Single pass numerically stable algorithm, see:
+        // https://en.wikipedia.org/wiki/Algorithms_for_calculating_variance#Online_algorithm
+        //
+        // n = 0, mean = 0.0, M2 = 0.0
+        //
+        // for x in data:
+        //     n += 1
+        //     delta = x - mean
+        //     mean += delta/n
+        //     delta2 = x - mean
+        //     M2 += delta*delta2
+        //
+        // if n < 2:
+        //     return float('nan')
+        // else:
+        //     return M2 / (n - 1)
+
+        // Register to store the number of elements read so far.
+        float count = 0.f, mean[ELEMENTS_PER_LDG], m2[ELEMENTS_PER_LDG];
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            mean[i] = 0.f;
+            m2[i] = 0.f;
+        }
+
+        // The number of elements loaded by this CTA.
+        int cta_count = 0;
+        // The base pointer to load from.
+        const uint16_t *gmem_src = &params.gmem_src[thread_c];
+
+        // outer loops
+        int OUTER_LOOPS = OUTER_LOOPS_ == 1? 1 : params.outer_loops;
+        // Load the batch of elements. Compute the mean/var across those elements.
+        const int pixels_per_iteration = PIXELS_PER_CTA_IN_REGISTERS*gridDim.x;
+
+        if (OUTER_LOOPS_ != 1) {
+            // We cannot load everything to store persistently, so let's makes sure registers and
+            // smem are fully utilized, offset is evenly divisible by 32
+            int offset = (pixels_per_iteration * OUTER_LOOPS +
+                          PIXELS_PER_CTA_IN_SMEM * gridDim.x - params.nhw) & ~31;
+            cta_nhw_regs -= offset;
+            cta_nhw_smem -= offset;
+        }
+
+        #pragma unroll 1
+        for (int loop_i = 0; loop_i < OUTER_LOOPS; ++loop_i) {
+            // The nhw position.
+            int nhw_regs = cta_nhw_regs + loop_i*pixels_per_iteration;
+            // Update the number of elements loaded by this CTA. TODO: Skip if <= 0!!!
+            cta_count += max(min(nhw_regs + PIXELS_PER_CTA_IN_REGISTERS, params.nhw) -
+                                 max(nhw_regs, 0), 0);
+
+            // Load the data and compute the local mean/sum and the variance.
+            if (USE_ONLINE_APPROACH) {
+                // Read the elements from memory.
+                float is_valid[PIXELS_PER_THREAD_IN_REGISTERS];
+                #pragma unroll
+                for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                    const int idx = nhw_regs + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                    zero_array(x_storage[i]);
+                    is_valid[i] = 0.f;
+                    if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {
+                        if (loop_i == OUTER_LOOPS - 1) {
+                            ldg_stream(x_storage[i], &gmem_src[idx*params.c]);
+                        } else {
+                            ldg(x_storage[i], &gmem_src[idx*params.c]);
+                        }
+                        is_valid[i] = 1.f;
+                    }
+                }
+
+                // Do the math.
+                #pragma unroll
+                for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                    // Convert to float.
+                    float x_math[ELEMENTS_PER_LDG];
+                    to_float(x_math, x_storage[i]);
+
+                    // Update the count.
+                    count += is_valid[i];
+                    // Invert the count.
+                    float inv_count = is_valid[i] ? 1.f / count : 0.f;
+
+                    // Update the mean and m2 using deltas.
+                    #pragma unroll
+                    for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {
+                        float delta0 = x_math[j] - mean[j];
+                        mean[j] += delta0 * inv_count;
+                        float delta1 = x_math[j] - mean[j];
+                        m2[j] += delta0 * delta1 * is_valid[i];
+                    }
+                }
+            } else {
+                // Read the elements from memory.
+                #pragma unroll
+                for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                    const int idx = nhw_regs + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                    zero_array(x_storage[i]);
+                    if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {
+                        if (loop_i == OUTER_LOOPS - 1) {
+                            ldg_stream(x_storage[i], &gmem_src[idx*params.c]);
+                        } else {
+                            ldg(x_storage[i], &gmem_src[idx*params.c]);
+                        }
+                        count += 1.f;
+                    }
+                }
+
+                // Sum the elements in registers.
+                #pragma unroll
+                for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                    // Convert to float.
+                    float x_math[ELEMENTS_PER_LDG];
+                    to_float(x_math, x_storage[i]);
+
+                    // Update the mean and m2 using deltas.
+                    #pragma unroll
+                    for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {
+                        mean[j] += x_math[j];
+                    }
+                }
+
+                // Compute the mean.
+                float inv_count = 1.f / count;
+                #pragma unroll
+                for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {
+                    mean[j] *= inv_count;
+                }
+
+                // Compute the variance.
+                #pragma unroll
+                for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                    // Convert to float.
+                    float x_math[ELEMENTS_PER_LDG];
+                    to_float(x_math, x_storage[i]);
+
+                    // Is it a valid pixel?
+                    float is_valid = i < static_cast<int>(count) ? 1.f : 0.f;
+                    // Update the mean and m2 using deltas.
+                    #pragma unroll
+                    for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {
+                        m2[j] += (x_math[j] - mean[j]) * (x_math[j] - mean[j]) * is_valid;
+                    }
+                }
+            }
+        }
+
+        // The elements to load and store in SMEM.
+        int smem_nhw = OUTER_LOOPS*pixels_per_iteration + cta_nhw_smem;
+        // Load elements from SMEM, update the CTA count.
+        int pixels_in_smem = min(smem_nhw + PIXELS_PER_CTA_IN_SMEM, params.nhw) - max(smem_nhw, 0);
+        if (pixels_in_smem > 0) {
+            cta_count += pixels_in_smem;
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {
+                const int idx = smem_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                float is_pixel_valid = (((unsigned int)idx <
+                                         (unsigned int)params.nhw) && is_valid_c) ? 1.f : 0.f;
+
+                PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG];
+                ldg_stream(x_storage_local, &gmem_src[(is_pixel_valid ? idx : 0)*params.c]);
+
+                // The offset to store in SMEM.
+                const int offset = i*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                // Store in SMEM.
+                write_to_smem(&smem_storage_packed[offset], threadIdx.x, x_storage_local);
+                // Update the count.
+                count += is_pixel_valid;
+                // Invert the count.
+                float inv_count = is_pixel_valid ? 1.f / count : 0.f;
+
+                float x_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage_local);
+                // Update the mean and m2 using deltas.
+                #pragma unroll
+                for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {
+                    float delta0 = x_math[j] - mean[j];
+                    mean[j] += delta0 * inv_count;
+                    float delta1 = x_math[j] - mean[j];
+                    m2[j] += delta0 * delta1 * is_pixel_valid;
+                }
+            }
+        }
+
+        // We scale the mean by the number of elements. It brings more stability.
+        float m1[ELEMENTS_PER_LDG];
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            m1[i] = mean[i] * count;
+        }
+
+        // Run the parallel sum accross the CTA to get the local sum.
+        ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+            smem, m1, thread_in_cta_nhw);
+        __syncthreads();
+
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(m1, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // Adjust the variance.
+        float inv_cta_count = 1.f / static_cast<float>(cta_count);
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            float mean_diff = m1[i]*inv_cta_count - mean[i];
+            m2[i] = m2[i] + mean_diff * mean_diff * count;
+        }
+
+        // Run the parallel sum accross the CTA to get the local adjusted variance.
+        ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+            smem, m2, thread_in_cta_nhw);
+
+        // The workspace in global memory is distributed across the different CTA.
+        int gmem_sums_offset = c_blk_index*gridDim.x*C_ELEMENTS_PER_CTA*2;
+
+        // Write the data for the CTA to global memory.
+        float *gmem_sums = &params.gmem_sums[gmem_sums_offset];
+        if (threadIdx.x < THREADS_PER_PIXEL) {
+            const int idx = blockIdx.x*THREADS_PER_PIXEL + threadIdx.x;
+            write_to_gmem(&gmem_sums[                           0], idx, m1);
+            write_to_gmem(&gmem_sums[C_ELEMENTS_PER_CTA*gridDim.x], idx, m2);
+        }
+
+        // The memory location to store the number of pixels per CTA.
+        int *gmem_counts = &params.gmem_counts[c_blk_index*gridDim.x];
+        if (threadIdx.x == 0) {
+            gmem_counts[blockIdx.x] = cta_count;
+        }
+
+        // Read the bias and scale.
+        float bias[ELEMENTS_PER_LDG], scale[ELEMENTS_PER_LDG];
+        if (is_valid_c) {
+            read_from_gmem(bias, &params.gmem_bias[cta_c], thread_in_cta_c);
+            read_from_gmem(scale, &params.gmem_scale[cta_c], thread_in_cta_c);
+        }
+
+        // The counters to count how many CTAs have retired at this point.
+        // A given cta uses the same counter every other time through the outer loop.
+        int *gmem_retired_ctas = &params.gmem_retired_ctas[c_blk_index % (2 * gridDim.y)];
+        inter_block_sync(gmem_retired_ctas, gridDim.x, blockIdx.x == 0);
+
+        // Reset the mean to compute the global mean.
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            m1[i] = 0.f;
+        }
+
+        // Build the global mean.
+        #pragma unroll 1
+        for (int idx = threadIdx.x; idx < THREADS_PER_PIXEL*gridDim.x; idx += THREADS_PER_CTA) {
+            float tmp[ELEMENTS_PER_LDG];
+            read_from_gmem(tmp, gmem_sums, idx);
+            add(m1, tmp);
+        }
+
+        if (params.sync_iters>0)
+        {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatchX<THREADS_PER_CTA>(
+                smem, m1, thread_in_cta_nhw, params.my_data, params.pair_datas, 4*c_blk_index+3, params.magic, params.sync_iters);
+        } else {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+                smem, m1, thread_in_cta_nhw);
+        }
+        __syncthreads();
+
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(m1, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // Normalize the mean.
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            m1[i] = m1[i] * params.svar_inv_count;
+        }
+
+        // Reset the variance.
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            m2[i] = 0.f;
+        }
+
+        // for add+relu fusion
+        const uint16_t *gmem_src1 = nullptr;
+        if (USE_ADD_RELU) {
+            gmem_src1 = &params.gmem_src1[thread_c];
+        }
+
+        // Build the global variance.
+        #pragma unroll 1
+        for (int idx = threadIdx.x; idx < THREADS_PER_PIXEL*gridDim.x; idx += THREADS_PER_CTA) {
+            // Read the means computed by different CTAs (again). Reuse tmp if we have 1 iteration.
+            float tmp_mean[ELEMENTS_PER_LDG], tmp_var[ELEMENTS_PER_LDG];
+            read_from_gmem(tmp_mean, &gmem_sums[                           0], idx);
+            read_from_gmem(tmp_var,  &gmem_sums[C_ELEMENTS_PER_CTA*gridDim.x], idx);
+
+            // Read the number of pixels visited by a given CTA.
+            cta_count = __ldg(&gmem_counts[idx / THREADS_PER_PIXEL]);
+
+            // Compute the diff to update the variance.
+            float mean_diff[ELEMENTS_PER_LDG], inv_cta_count = 1.f / static_cast<float>(cta_count);
+            #pragma unroll
+            for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+                mean_diff[i] = m1[i] - tmp_mean[i]*inv_cta_count;
+            }
+
+            // Update the variance.
+            #pragma unroll
+            for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+                m2[i] += tmp_var[i] + mean_diff[i]*mean_diff[i]*static_cast<float>(cta_count);
+            }
+        }
+
+        if (params.sync_iters>0)
+        {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatchX<THREADS_PER_CTA>(
+                smem, m2, thread_in_cta_nhw, params.my_data, params.pair_datas, 4*c_blk_index+2, params.magic, params.sync_iters);
+        } else {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+                smem, m2, thread_in_cta_nhw);
+        }
+        __syncthreads();
+
+        read_from_smem(m2, smem, thread_in_cta_c);
+
+        // Finalize the stddev.
+        // becasue saved var and running var may have different denominator, we don't do it here
+        // scale_(m2, inv_count);
+
+        // store the saved mean/var
+        float svarinv[ELEMENTS_PER_LDG];
+        bool is_valid_for_saving = is_valid_c && blockIdx.x == 0 && thread_in_cta_nhw == 0;
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            svarinv[i] = rsqrtf(m2[i] * params.svar_inv_count + params.var_eps);
+        }
+        if (is_valid_for_saving) {
+            write_to_gmem(params.gmem_saved_mean, thread_c/ELEMENTS_PER_LDG, m1);
+            write_to_gmem(params.gmem_saved_var, thread_c/ELEMENTS_PER_LDG, svarinv);
+        }
+
+        // store the running mean/var
+        float rmean[ELEMENTS_PER_LDG], rvar[ELEMENTS_PER_LDG];
+        zero_array(rmean);
+        zero_array(rvar);
+        if (params.exp_avg_factor != 1.f && is_valid_for_saving) {
+            read_from_gmem(rmean, params.gmem_running_mean, thread_c/ELEMENTS_PER_LDG);
+            read_from_gmem(rvar, params.gmem_running_var, thread_c/ELEMENTS_PER_LDG);
+        }
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            rmean[i] = (1.f - params.exp_avg_factor) * rmean[i] +   \
+                params.exp_avg_factor * m1[i];
+            rvar[i] = (1.f - params.exp_avg_factor) * rvar[i] +     \
+                params.exp_avg_factor * (m2[i] * params.rvar_inv_count);
+        }
+        if (is_valid_for_saving) {
+            write_to_gmem(params.gmem_running_mean, thread_c/ELEMENTS_PER_LDG, rmean);
+            write_to_gmem(params.gmem_running_var, thread_c/ELEMENTS_PER_LDG, rvar);
+        }
+
+        // Update the scale with the stddev and eps.
+        multiply(scale, svarinv);
+
+        // The base pointer to write to.
+        uint16_t *const gmem_dst = &params.gmem_dst[thread_c];
+
+        unsigned int *const gmem_relu_bitmask = params.gmem_relu_bitmask +
+                                     ((params.nhw + 31) & ~31) * 2 * c_blk_index;
+
+        // Store the elements in registers.
+        #pragma unroll 1
+        for (int loop_i = OUTER_LOOPS-1; loop_i >= 0; --loop_i) {
+            // The value for nhw.
+            int out_nhw = cta_nhw_regs + loop_i*pixels_per_iteration;
+
+            // Normalize the elements and write to memory.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                const int idx = out_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                const bool is_valid_nhw =
+                    static_cast<unsigned int>(idx) < static_cast<unsigned int>(params.nhw);
+                const bool is_valid = is_valid_nhw && is_valid_c;
+                // Convert to float.
+                float x_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage[i]);
+
+                // Normalize and apply activation function
+                normalize(x_math, bias, scale, m1);
+                if (USE_ADD_RELU) {
+                    float x1_math[ELEMENTS_PER_LDG];
+                    ldg_stream(x1_math, &gmem_src1[(is_valid ? idx : 0)*params.c]);
+                    add(x_math, x1_math);
+                    unsigned int relu_mask;
+                    int lane_id = threadIdx.x & 31;
+                    #pragma unroll
+                    for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+                        bool rectified = x_math[i] < 0.0F;
+                        unsigned int local_relu_mask = __ballot_sync(0xFFFFFFFFU, rectified);
+                        if (lane_id == i) {
+                            // Thread 0 remembers the relu_mask from the first time through this
+                            // loop, Thread 1 the next, Thread 2 the next, and Thread 3 the last.
+                            relu_mask = local_relu_mask;
+                        }
+                        if (rectified) {
+                            x_math[i] = 0.0F;
+                        }
+                    }
+                    if (is_valid_nhw && (lane_id < ELEMENTS_PER_LDG)) {
+                        gmem_relu_bitmask[idx * 2 + lane_id] = relu_mask;
+                    }
+                } else if (USE_RELU) {
+                    relu_activation(x_math);
+                }
+
+                // Write back.
+                if (is_valid) {
+                    stg_stream(&gmem_dst[idx*params.c], x_math);
+                }
+            }
+
+            // The next value of nhw.
+            out_nhw -= pixels_per_iteration;
+
+            // Read the next elements from memory.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                const int idx = out_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {
+                    ldg_stream(x_storage[i], &gmem_src[idx*params.c]);
+                }
+            }
+        }
+
+        // Normalize the elements from SMEM and write them out.
+        if (pixels_in_smem > 0) {
+            #pragma unroll 2
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {
+                const int idx = smem_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                const bool is_valid_nhw =
+                    static_cast<unsigned int>(idx) < static_cast<unsigned int>(params.nhw);
+                const bool is_valid = is_valid_nhw && is_valid_c;
+
+                // Read from SMEM.
+                const int offset = i*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG];
+                read_from_smem(x_storage_local, &smem_storage_packed[offset], threadIdx.x);
+                float x_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage_local);
+
+                // Normalize and apply activation function
+                normalize(x_math, bias, scale, m1);
+                if (USE_ADD_RELU) {
+                    float x1_math[ELEMENTS_PER_LDG];
+                    ldg_stream(x1_math, &gmem_src1[(is_valid ? idx : 0)*params.c]);
+                    add(x_math, x1_math);
+                    unsigned int relu_mask;
+                    int lane_id = threadIdx.x & 31;
+                    #pragma unroll
+                    for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+                        bool rectified = x_math[i] < 0.0F;
+                        unsigned int local_relu_mask = __ballot_sync(0xFFFFFFFFU, rectified);
+                        if (lane_id == i) {
+                            relu_mask = local_relu_mask;
+                        }
+                        if (rectified) {
+                            x_math[i] = 0.0F;
+                        }
+                    }
+                    if (is_valid_nhw && (lane_id < ELEMENTS_PER_LDG)) {
+                        gmem_relu_bitmask[idx * 2 + lane_id] = relu_mask;
+                    }
+                } else if (USE_RELU) {
+                    relu_activation(x_math);
+                }
+
+                // Write back.
+                if (is_valid) {
+                    stg_stream(&gmem_dst[idx*params.c], x_math);
+                }
+            }
+        }
+        // We're about to start on the next c-blk.  Needed?
+        __syncthreads();
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct NhwcBatchNormBwdParams {
+    // The input/output tensors.
+    uint16_t *gmem_src, *gmem_dy, *gmem_dst, *gmem_dst1;
+    // dscale/dbias
+    float *gmem_dscale, *gmem_dbias;
+    // The scale and bias.
+    float *gmem_scale, *gmem_bias;
+    // The mean/inv-var saved from fwd pass
+    float *gmem_saved_mean, *gmem_saved_var;
+    // ReLU bitmask
+    unsigned int *gmem_relu_bitmask;
+    // The dimensions.
+    int nhw, c;
+    // factor to scale sum of squared errors to get saved variance.  Must be 1/nhw.
+    float svar_inv_count;
+    // The buffer to do the reduction for dscale and dbias
+    float *gmem_sums;
+    // The counters of retired CTAs.
+    int *gmem_retired_ctas;
+    // outer loop count
+    int outer_loops;
+    // number of CTAs along .x dimension
+    int c_blks;
+
+    void* my_data;
+    void* pair_datas[4];
+    int magic;
+    int sync_iters;
+    float wgrad_coeff;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void relu_bwd(float (&dy)[N], const float (&x)[N],
+                              const float (&mean_var_scale_bias)[N],
+                              const float (&var_scale)[N], bool valid_data) {
+    #pragma unroll
+    for (int j = 0; j < N; ++j) {
+        float y = (x[j] * var_scale[j]) + mean_var_scale_bias[j];
+        if ((y <= 0.f) && valid_data) {
+            dy[j] = 0.f;
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void relu_bwd(float (&dy)[N], const float (&y)[N], bool valid_data) {
+    #pragma unroll
+    for (int j = 0; j < N; ++j) {
+        if ((y[j] <= 0.f) && valid_data) {
+            dy[j] = 0.f;
+        }
+    }
+}
+
+template <int N>
+DEVICE_FUNCTION void relu_bwd(float (&dy)[N], const bool (&rectified)[N], bool valid_data) {
+    #pragma unroll
+    for (int j = 0; j < N; ++j) {
+        if (rectified[j] && valid_data) {
+            dy[j] = 0.f;
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void relu_bwd_for_dx(float (&dy)[N],
+                                     const float (&x)[N],
+                                     const float (&mean_var_scale_bias)[N],
+                                     const float (&var_scale)[N]) {
+    #pragma unroll
+    for (int j = 0; j < N; ++j) {
+        float y = (x[j] * var_scale[j]) + mean_var_scale_bias[j];
+        if (y <= 0.f) {
+            dy[j] = 0.f;
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void relu_bwd_for_dx(float (&dy)[N], const float (&y)[N]) {
+    #pragma unroll
+    for (int j = 0; j < N; ++j) {
+        if (y[j] <= 0.f) {
+            dy[j] = 0.f;
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void bwd_update(float (&dscale)[N], float (&dbias)[N],
+                                const float (&dy)[N], const float (&x)[N],
+                                const float (&mean)[N], float inv_count) {
+    #pragma unroll
+    for (int j = 0; j < N; ++j) {
+        float delta0 = dy[j] - dbias[j];
+        dbias[j] += delta0 * inv_count;
+        delta0 = (dy[j] * (x[j] - mean[j])) - dscale[j];
+        dscale[j] += delta0 * inv_count;
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <int N>
+DEVICE_FUNCTION void bwd_dx(float (&dx)[N], const float (&dy)[N],
+                            const float (&var)[N], const float (&x)[N], const float (&mean)[N],
+                            const float (&dscale)[N], const float (&dbias)[N], float inv_count) {
+    #pragma unroll
+    for (int j = 0; j < N; ++j) {
+        float tmp1 = dy[j] - (dbias[j]* inv_count);
+        float tmp2 = dscale[j] * inv_count;
+        float tmp3 = x[j] - mean[j];
+        dx[j] = var[j] * (tmp1 - (tmp2 * tmp3));
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    typename Storage,
+    int THREADS_PER_CTA,
+    int THREADS_PER_PIXEL,
+    int PIXELS_PER_THREAD_IN_REGISTERS,
+    int PIXELS_PER_THREAD_IN_SMEM,
+    int ELEMENTS_PER_LDG,
+    int USE_ONLINE_APPROACH,
+    int OUTER_LOOPS_,
+    int DESIRED_OCCUPANCY
+>
+__global__ __launch_bounds__(THREADS_PER_CTA, DESIRED_OCCUPANCY)
+    void nhwc_batch_norm_bwd(NhwcBatchNormBwdParams params) {
+    // The number of pixels loaded in a single LDG.
+    const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;
+    // The number of pixels computed per CTA stored in registers.
+    const int PIXELS_PER_CTA_IN_REGISTERS = PIXELS_PER_THREAD_IN_REGISTERS * PIXELS_PER_LDG;
+    // The number of pixels computed per CTA stored in SMEM.
+    const int PIXELS_PER_CTA_IN_SMEM = PIXELS_PER_THREAD_IN_SMEM*PIXELS_PER_LDG;
+    // The number of C elements per CTA.
+    const int C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL*ELEMENTS_PER_LDG;
+
+    // Shared memory to do CTA-wide parallel sums.
+    __shared__ float smem[THREADS_PER_PIXEL*(THREADS_PER_CTA/32)*ELEMENTS_PER_LDG];
+
+    // The adapter for the storage.
+    typedef PackedStorage<Storage, ELEMENTS_PER_LDG> PackedStorage_;
+    // The data type for packed storage in SMEM.
+    typedef typename PackedStorage_::Type PackedStorageType;
+    // The number of elements in the packed storage.
+    const int PACKED_ELEMENTS_PER_LDG = PackedStorage_::PACKED_ELEMENTS_PER_LDG;
+    // Registers to keep the data live for the persistent approach.
+    PackedStorageType x_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];
+    PackedStorageType dy_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];
+
+    // Shared memory buffer to store the extra pixels.
+    extern __shared__ PackedStorageType smem_storage_packed[];
+
+    for (int c_blk_index = blockIdx.y; c_blk_index < params.c_blks; c_blk_index += gridDim.y) {
+        // The position in the NHW dimension where the CTA starts.
+        int cta_nhw_regs = blockIdx.x * PIXELS_PER_CTA_IN_REGISTERS;
+        // The position in the NHW dimension where the CTA starts for the portion in SMEM.
+        int cta_nhw_smem = blockIdx.x * PIXELS_PER_CTA_IN_SMEM;
+        // Compute the NHW coordinate of the thread in the CTA.
+        const int thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;
+
+        // The position in the C dimension where the CTA starts.
+        const int cta_c = c_blk_index * C_ELEMENTS_PER_CTA;
+        // Compute the C coordinate of the thread in the CTA.
+        const int thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;
+        // Compute the C coordinate of the thread.
+        const int thread_c = cta_c + thread_in_cta_c*ELEMENTS_PER_LDG;
+
+        // Is the thread working on a valid C dimension?
+        const int is_valid_c = thread_c < params.c;
+
+        // Registers to store the mean used for entire duration
+        float mean[ELEMENTS_PER_LDG];
+        zero_array(mean);
+        if (is_valid_c) {
+            read_from_gmem(mean, params.gmem_saved_mean, thread_c/ELEMENTS_PER_LDG);
+        }
+
+        // accumulation related registers
+        float count = 0.f, dscale[ELEMENTS_PER_LDG], dbias[ELEMENTS_PER_LDG];
+        zero_array(dscale);
+        zero_array(dbias);
+
+        // The number of elements loaded by this CTA.
+        int cta_count = 0;
+        // The base pointers to load from.
+        const uint16_t *gmem_src = &params.gmem_src[thread_c];
+        const uint16_t *gmem_dy = &params.gmem_dy[thread_c];
+
+        // outer loops
+        int OUTER_LOOPS = OUTER_LOOPS_ == 1? 1 : params.outer_loops;
+        // Load the batch of elements. Compute sum across them
+        const int pixels_per_iteration = PIXELS_PER_CTA_IN_REGISTERS*gridDim.x;
+
+        if (OUTER_LOOPS_ != 1) {
+            // We cannot load everything to store persistently, so let's makes sure registers and
+            // smem are fully utilized
+            int offset = params.nhw - pixels_per_iteration * OUTER_LOOPS -
+                         PIXELS_PER_CTA_IN_SMEM * gridDim.x;
+            cta_nhw_regs += offset;
+            cta_nhw_smem += offset;
+        }
+
+        #pragma unroll 1
+        for (int loop_i = 0; loop_i < OUTER_LOOPS; ++loop_i) {
+            // The nhw position.
+            int nhw_regs = cta_nhw_regs + loop_i*pixels_per_iteration;
+            // Update the number of elements loaded by this CTA. TODO: Skip if <= 0!!!
+            cta_count += max(0, min(PIXELS_PER_CTA_IN_REGISTERS, params.nhw-nhw_regs));
+
+            // Read the elements from memory.
+            float is_valid[PIXELS_PER_THREAD_IN_REGISTERS];
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                const int idx = nhw_regs + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                zero_array(x_storage[i]);
+                zero_array(dy_storage[i]);
+                is_valid[i] = 0.f;
+                if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {
+                    if (loop_i == OUTER_LOOPS - 1) {
+                        ldg_stream(x_storage[i], &gmem_src[idx*params.c]);
+                        ldg_stream(dy_storage[i], &gmem_dy[idx*params.c]);
+                    } else {
+                        ldg(x_storage[i], &gmem_src[idx*params.c]);
+                        ldg(dy_storage[i], &gmem_dy[idx*params.c]);
+                    }
+                    is_valid[i] = 1.f;
+                }
+            }
+
+            // Do the math.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                // Convert to float and update
+                float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage[i]);
+                to_float(dy_math, dy_storage[i]);
+
+                // Update the count.
+                count += is_valid[i];
+                // Invert the count.
+                float inv_count = is_valid[i] ? 1.f / count : 0.f;
+
+                bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);
+            }
+        }
+
+        // The elements to load and store in SMEM.
+        int smem_nhw = OUTER_LOOPS*pixels_per_iteration + cta_nhw_smem;
+        // Load elements from SMEM, update the CTA count.
+        int pixels_in_smem = min(PIXELS_PER_CTA_IN_SMEM, params.nhw-smem_nhw);
+        if (pixels_in_smem > 0) {
+            cta_count += pixels_in_smem;
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {
+                const int idx = smem_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                bool is_pixel_valid = (((unsigned int)idx <
+                                        (unsigned int)params.nhw) && is_valid_c);
+                PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG],
+                                  dy_storage_local[PACKED_ELEMENTS_PER_LDG];
+                zero_array(x_storage_local);
+                zero_array(dy_storage_local);
+                if (is_pixel_valid) {
+                    ldg_stream(x_storage_local, &gmem_src[idx*params.c]);
+                    ldg_stream(dy_storage_local, &gmem_dy[idx*params.c]);
+                }
+
+                // The offset to store in SMEM.
+                int offset = i*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                // Store in SMEM.
+                write_to_smem(&smem_storage_packed[offset], threadIdx.x, x_storage_local);
+                offset += PIXELS_PER_THREAD_IN_SMEM*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                write_to_smem(&smem_storage_packed[offset], threadIdx.x, dy_storage_local);
+                // Update the count.
+                count += is_pixel_valid;
+                // Invert the count.
+                float inv_count = is_pixel_valid ? 1.f / count : 0.f;
+
+                float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage_local);
+                to_float(dy_math, dy_storage_local);
+
+                bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);
+            }
+        }
+
+        // We scale the mean by the number of elements. It brings more stability.
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            dbias[i] *= count;
+            dscale[i] *= count;
+        }
+
+        // dscale parallel sum
+        ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+            smem, dscale, thread_in_cta_nhw);
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dscale, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // dbias parallel sum
+        ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+            smem, dbias, thread_in_cta_nhw);
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dbias, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // The workspace in global memory is distributed across the different CTA.
+        int gmem_sums_offset = c_blk_index*gridDim.x*C_ELEMENTS_PER_CTA*2;
+        // Write the data for the CTA to global memory.
+        float *gmem_sums = &params.gmem_sums[gmem_sums_offset];
+        if (threadIdx.x < THREADS_PER_PIXEL) {
+            const int idx = blockIdx.x*THREADS_PER_PIXEL + threadIdx.x;
+            write_to_gmem(&gmem_sums[                           0], idx, dscale);
+            write_to_gmem(&gmem_sums[C_ELEMENTS_PER_CTA*gridDim.x], idx, dbias);
+        }
+
+        // The counters to count how many CTAs have retired at this point.
+        // A given cta uses the same counter every other time through the outer loop.
+        int *gmem_retired_ctas = &params.gmem_retired_ctas[c_blk_index % (2 * gridDim.y)];
+        inter_block_sync(gmem_retired_ctas, gridDim.x, blockIdx.x == 0);
+
+        // Reset the accumulators for global summation
+        zero_array(dscale);
+        zero_array(dbias);
+
+        // Build the global accumulation
+        #pragma unroll 1
+        for (int idx = threadIdx.x; idx < THREADS_PER_PIXEL*gridDim.x; idx += THREADS_PER_CTA) {
+            float tmp1[ELEMENTS_PER_LDG], tmp2[ELEMENTS_PER_LDG];
+            read_from_gmem(tmp1, gmem_sums,                              idx);
+            read_from_gmem(tmp2, gmem_sums+C_ELEMENTS_PER_CTA*gridDim.x, idx);
+
+            #pragma unroll
+            for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+                dscale[i] += tmp1[i];
+                dbias[i] += tmp2[i];
+            }
+        }
+
+        // dscale parallel sum
+        if (params.sync_iters>0) {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatchX<THREADS_PER_CTA>(
+                smem, dscale, thread_in_cta_nhw, params.my_data, params.pair_datas, 4*c_blk_index+1, params.magic, params.sync_iters);
+        } else {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+                smem, dscale, thread_in_cta_nhw);
+        }
+
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dscale, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // dbias parallel sum
+        if (params.sync_iters>0) {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatchX<THREADS_PER_CTA>(
+                smem, dbias, thread_in_cta_nhw, params.my_data, params.pair_datas, 4*c_blk_index+0, params.magic, params.sync_iters);
+        } else {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+                smem, dbias, thread_in_cta_nhw);
+        }
+
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dbias, smem, thread_in_cta_c);
+
+        // inv-var
+        float var[ELEMENTS_PER_LDG];
+        zero_array(var);
+        if (is_valid_c) {
+            read_from_gmem(var, params.gmem_saved_var, thread_c/ELEMENTS_PER_LDG);
+        }
+
+        // Normalize the dscale.
+        multiply(dscale, var);
+
+        // store dscale/dbias
+        bool is_valid_for_saving = is_valid_c && blockIdx.x == 0 && thread_in_cta_nhw == 0;
+        if (is_valid_for_saving) {
+            if (params.sync_iters>0)
+            {
+                scaled_write_to_gmem(params.gmem_dscale, thread_c/ELEMENTS_PER_LDG, dscale, params.wgrad_coeff);
+                scaled_write_to_gmem(params.gmem_dbias, thread_c/ELEMENTS_PER_LDG, dbias, params.wgrad_coeff);
+            } else {
+                write_to_gmem(params.gmem_dscale, thread_c/ELEMENTS_PER_LDG, dscale);
+                write_to_gmem(params.gmem_dbias, thread_c/ELEMENTS_PER_LDG, dbias);
+            }
+        }
+
+        // scale
+        float scale[ELEMENTS_PER_LDG];
+        zero_array(scale);
+        if (is_valid_c) {
+            read_from_gmem(scale, params.gmem_scale, thread_c/ELEMENTS_PER_LDG);
+        }
+
+        // Further normalize the dscale to be used in dx calculation
+        multiply(dscale, var);
+        // scale the inv-var as well, afterwards
+        multiply(var, scale);
+
+        // inverse count
+        float inv_count = params.svar_inv_count;
+
+        // The base pointer to write to.
+        uint16_t *const gmem_dst = &params.gmem_dst[thread_c];
+
+        // Store the elements in registers.
+        #pragma unroll 1
+        for (int loop_i = OUTER_LOOPS-1; loop_i >= 0; --loop_i) {
+            // The value for nhw.
+            int out_nhw = cta_nhw_regs + loop_i*pixels_per_iteration;
+
+            // Normalize the elements and write to memory.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                // Convert to float.
+                float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage[i]);
+                to_float(dy_math, dy_storage[i]);
+
+                float dx[ELEMENTS_PER_LDG];
+                bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);
+
+                // Write back.
+                const int idx = out_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {
+                    stg_stream(&gmem_dst[idx*params.c], dx);
+                }
+            }
+
+            // The next value of nhw.
+            out_nhw -= pixels_per_iteration;
+
+            // Read the next elements from memory.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                const int idx = out_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {
+                    ldg_stream(x_storage[i], &gmem_src[idx*params.c]);
+                    ldg_stream(dy_storage[i], &gmem_dy[idx*params.c]);
+                }
+            }
+        }
+
+        // Normalize the elements from SMEM and write them out.
+        if (pixels_in_smem > 0) {
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {
+                const int idx = smem_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                const bool is_valid = ((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c;
+                if (is_valid) {
+                    // Read from SMEM.
+                    int offset = i*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                    PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG],
+                        dy_storage_local[PACKED_ELEMENTS_PER_LDG];
+                    read_from_smem(x_storage_local, &smem_storage_packed[offset], threadIdx.x);
+                    offset += PIXELS_PER_THREAD_IN_SMEM*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                    read_from_smem(dy_storage_local, &smem_storage_packed[offset], threadIdx.x);
+                    float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                    to_float(x_math, x_storage_local);
+                    to_float(dy_math, dy_storage_local);
+
+                    float dx[ELEMENTS_PER_LDG];
+                    bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);
+
+                    // Write back.
+                    stg_stream(&gmem_dst[idx*params.c], dx);
+                }
+            }
+        }
+        // We're about to start on the next c-blk.  Needed?
+        __syncthreads();
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    typename Storage,
+    int THREADS_PER_CTA,
+    int THREADS_PER_PIXEL,
+    int PIXELS_PER_THREAD_IN_REGISTERS,
+    int PIXELS_PER_THREAD_IN_SMEM,
+    int ELEMENTS_PER_LDG,
+    int USE_ONLINE_APPROACH,
+    int OUTER_LOOPS_,
+    int DESIRED_OCCUPANCY
+>
+__global__ __launch_bounds__(THREADS_PER_CTA, DESIRED_OCCUPANCY)
+    void nhwc_batch_norm_bwd_relu(NhwcBatchNormBwdParams params) {
+    // The number of pixels loaded in a single LDG.
+    const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;
+    // The number of pixels computed per CTA stored in registers.
+    const int PIXELS_PER_CTA_IN_REGISTERS = PIXELS_PER_THREAD_IN_REGISTERS * PIXELS_PER_LDG;
+    // The number of pixels computed per CTA stored in SMEM.
+    const int PIXELS_PER_CTA_IN_SMEM = PIXELS_PER_THREAD_IN_SMEM*PIXELS_PER_LDG;
+    // The number of C elements per CTA.
+    const int C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL*ELEMENTS_PER_LDG;
+
+    // Shared memory to do CTA-wide parallel sums.
+    __shared__ float smem[THREADS_PER_PIXEL*(THREADS_PER_CTA/32)*ELEMENTS_PER_LDG];
+
+    // The adapter for the storage.
+    typedef PackedStorage<Storage, ELEMENTS_PER_LDG> PackedStorage_;
+    // The data type for packed storage in SMEM.
+    typedef typename PackedStorage_::Type PackedStorageType;
+    // The number of elements in the packed storage.
+    const int PACKED_ELEMENTS_PER_LDG = PackedStorage_::PACKED_ELEMENTS_PER_LDG;
+    // Registers to keep the data live for the persistent approach.
+    PackedStorageType x_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];
+    PackedStorageType dy_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];
+
+    // Shared memory buffer to store the extra pixels.
+    extern __shared__ PackedStorageType smem_storage_packed[];
+
+    for (int c_blk_index = blockIdx.y; c_blk_index < params.c_blks; c_blk_index += gridDim.y) {
+        // The position in the NHW dimension where the CTA starts.
+        int cta_nhw_regs = blockIdx.x * PIXELS_PER_CTA_IN_REGISTERS;
+        // The position in the NHW dimension where the CTA starts for the portion in SMEM.
+        int cta_nhw_smem = blockIdx.x * PIXELS_PER_CTA_IN_SMEM;
+        // Compute the NHW coordinate of the thread in the CTA.
+        const int thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;
+
+        // The position in the C dimension where the CTA starts.
+        const int cta_c = c_blk_index * C_ELEMENTS_PER_CTA;
+        // Compute the C coordinate of the thread in the CTA.
+        const int thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;
+        // Compute the C coordinate of the thread.
+        const int thread_c = cta_c + thread_in_cta_c*ELEMENTS_PER_LDG;
+
+        // Is the thread working on a valid C dimension?
+        const int is_valid_c = thread_c < params.c;
+
+
+        // Registers to store the mean/var/scale/bias used for the entire duration
+        // Register usage optimizations:
+        // 1. Can combine bias - (mean * var * scale) into a single register
+        // 2. Can combine var * scale into a single register
+        float varscale[ELEMENTS_PER_LDG];
+        zero_array(varscale);
+        if (is_valid_c) {
+            read_from_gmem(varscale, params.gmem_saved_var, thread_c/ELEMENTS_PER_LDG);
+        }
+        float tmp[ELEMENTS_PER_LDG];
+        zero_array(tmp);
+        if (is_valid_c) {
+            read_from_gmem(tmp, params.gmem_scale, thread_c/ELEMENTS_PER_LDG);
+        }
+        multiply(varscale, tmp);
+        float mean[ELEMENTS_PER_LDG];
+        zero_array(mean);
+        if (is_valid_c) {
+            read_from_gmem(mean, params.gmem_saved_mean, thread_c/ELEMENTS_PER_LDG);
+        }
+        zero_array(tmp);
+        if (is_valid_c) {
+            read_from_gmem(tmp, params.gmem_bias, thread_c/ELEMENTS_PER_LDG);
+        }
+        float mean_var_scale_bias[ELEMENTS_PER_LDG];
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            mean_var_scale_bias[i] = tmp[i] - (mean[i] * varscale[i]);
+        }
+
+        // accumulation related registers
+        float count = 0.f, dscale[ELEMENTS_PER_LDG], dbias[ELEMENTS_PER_LDG];
+        zero_array(dscale);
+        zero_array(dbias);
+
+        // The number of elements loaded by this CTA.
+        int cta_count = 0;
+        // The base pointers to load from.
+        const uint16_t *gmem_src = &params.gmem_src[thread_c];
+        const uint16_t *gmem_dy = &params.gmem_dy[thread_c];
+
+        // outer loops
+        int OUTER_LOOPS = OUTER_LOOPS_ == 1? 1 : params.outer_loops;
+        // Load the batch of elements. Compute sum across them
+        const int pixels_per_iteration = PIXELS_PER_CTA_IN_REGISTERS*gridDim.x;
+
+        if (OUTER_LOOPS_ != 1) {
+            // We cannot load everything to store persistently, so let's makes sure registers and
+            // smem are fully utilized
+            int offset = params.nhw - pixels_per_iteration * OUTER_LOOPS -
+                         PIXELS_PER_CTA_IN_SMEM * gridDim.x;
+            cta_nhw_regs += offset;
+            cta_nhw_smem += offset;
+        }
+
+        #pragma unroll 1
+        for (int loop_i = 0; loop_i < OUTER_LOOPS; ++loop_i) {
+            // The nhw position.
+            int nhw_regs = cta_nhw_regs + loop_i*pixels_per_iteration;
+            // Update the number of elements loaded by this CTA. TODO: Skip if <= 0!!!
+            cta_count += max(0, min(PIXELS_PER_CTA_IN_REGISTERS, params.nhw-nhw_regs));
+
+            // Read the elements from memory.
+            float is_valid[PIXELS_PER_THREAD_IN_REGISTERS];
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                const int idx = nhw_regs + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                zero_array(x_storage[i]);
+                zero_array(dy_storage[i]);
+                is_valid[i] = 0.f;
+                if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {
+                    if (loop_i == OUTER_LOOPS - 1) {
+                        ldg_stream(x_storage[i], &gmem_src[idx*params.c]);
+                        ldg_stream(dy_storage[i], &gmem_dy[idx*params.c]);
+                    } else {
+                        ldg(x_storage[i], &gmem_src[idx*params.c]);
+                        ldg(dy_storage[i], &gmem_dy[idx*params.c]);
+                    }
+                    is_valid[i] = 1.f;
+                }
+            }
+
+            // Do the math.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                // Convert to float and update
+                float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage[i]);
+                to_float(dy_math, dy_storage[i]);
+
+                // Update the count.
+                count += is_valid[i];
+                // Invert the count.
+                float inv_count = is_valid[i] ? 1.f / count : 0.f;
+
+                relu_bwd(dy_math, x_math, mean_var_scale_bias, varscale, is_valid[i]);
+                bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);
+            }
+        }
+
+        // The elements to load and store in SMEM.
+        int smem_nhw = OUTER_LOOPS*pixels_per_iteration + cta_nhw_smem;
+        // Load elements from SMEM, update the CTA count.
+        int pixels_in_smem = min(PIXELS_PER_CTA_IN_SMEM, params.nhw-smem_nhw);
+        if (pixels_in_smem > 0) {
+            cta_count += pixels_in_smem;
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {
+                const int idx = smem_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                bool is_pixel_valid = (((unsigned int)idx <
+                                        (unsigned int)params.nhw) && is_valid_c);
+                PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG],
+                                  dy_storage_local[PACKED_ELEMENTS_PER_LDG];
+                zero_array(x_storage_local);
+                zero_array(dy_storage_local);
+                if (is_pixel_valid) {
+                    ldg_stream(x_storage_local, &gmem_src[idx*params.c]);
+                    ldg_stream(dy_storage_local, &gmem_dy[idx*params.c]);
+                }
+
+                // The offset to store in SMEM.
+                int offset = i*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                // Store in SMEM.
+                write_to_smem(&smem_storage_packed[offset], threadIdx.x, x_storage_local);
+                offset += PIXELS_PER_THREAD_IN_SMEM*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                write_to_smem(&smem_storage_packed[offset], threadIdx.x, dy_storage_local);
+                // Update the count.
+                count += is_pixel_valid;
+                // Invert the count.
+                float inv_count = is_pixel_valid ? 1.f / count : 0.f;
+
+                float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage_local);
+                to_float(dy_math, dy_storage_local);
+
+                relu_bwd(dy_math, x_math, mean_var_scale_bias, varscale, is_pixel_valid);
+                bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);
+            }
+        }
+
+        // We scale the mean by the number of elements. It brings more stability.
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            dbias[i] *= count;
+            dscale[i] *= count;
+        }
+
+        // dscale parallel sum
+        ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+            smem, dscale, thread_in_cta_nhw);
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dscale, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // dbias parallel sum
+        ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+            smem, dbias, thread_in_cta_nhw);
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dbias, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // The workspace in global memory is distributed across the different CTA.
+        int gmem_sums_offset = c_blk_index*gridDim.x*C_ELEMENTS_PER_CTA*2;
+        // Write the data for the CTA to global memory.
+        float *gmem_sums = &params.gmem_sums[gmem_sums_offset];
+        if (threadIdx.x < THREADS_PER_PIXEL) {
+            const int idx = blockIdx.x*THREADS_PER_PIXEL + threadIdx.x;
+            write_to_gmem(&gmem_sums[                           0], idx, dscale);
+            write_to_gmem(&gmem_sums[C_ELEMENTS_PER_CTA*gridDim.x], idx, dbias);
+        }
+
+        // The counters to count how many CTAs have retired at this point.
+        // A given cta uses the same counter every other time through the outer loop.
+        int *gmem_retired_ctas = &params.gmem_retired_ctas[c_blk_index % (2 * gridDim.y)];
+        inter_block_sync(gmem_retired_ctas, gridDim.x, blockIdx.x == 0);
+
+        // Reset the accumulators for global summation
+        zero_array(dscale);
+        zero_array(dbias);
+
+        // Build the global accumulation
+        #pragma unroll 1
+        for (int idx = threadIdx.x; idx < THREADS_PER_PIXEL*gridDim.x; idx += THREADS_PER_CTA) {
+            float tmp1[ELEMENTS_PER_LDG], tmp2[ELEMENTS_PER_LDG];
+            read_from_gmem(tmp1, gmem_sums,                              idx);
+            read_from_gmem(tmp2, gmem_sums+C_ELEMENTS_PER_CTA*gridDim.x, idx);
+
+            #pragma unroll
+            for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+                dscale[i] += tmp1[i];
+                dbias[i] += tmp2[i];
+            }
+        }
+
+        // dscale parallel sum
+        if (params.sync_iters>0) {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatchX<THREADS_PER_CTA>(
+                smem, dscale, thread_in_cta_nhw, params.my_data, params.pair_datas, 4*c_blk_index+1, params.magic, params.sync_iters);
+        } else {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+                smem, dscale, thread_in_cta_nhw);
+        }
+
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dscale, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // dbias parallel sum
+        if (params.sync_iters>0) {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatchX<THREADS_PER_CTA>(
+                smem, dbias, thread_in_cta_nhw, params.my_data, params.pair_datas, 4*c_blk_index+0, params.magic, params.sync_iters);
+        } else {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+                smem, dbias, thread_in_cta_nhw);
+        }
+
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dbias, smem, thread_in_cta_c);
+
+        // Normalize the dscale.
+        float var[ELEMENTS_PER_LDG];
+        zero_array(var);
+        if (is_valid_c) {
+            read_from_gmem(var, params.gmem_saved_var, thread_c/ELEMENTS_PER_LDG);
+        }
+        multiply(dscale, var);
+
+        // store dscale/dbias
+        bool is_valid_for_saving = is_valid_c && blockIdx.x == 0 && thread_in_cta_nhw == 0;
+        if (is_valid_for_saving) {
+            if (params.sync_iters>0)
+            {
+                scaled_write_to_gmem(params.gmem_dscale, thread_c/ELEMENTS_PER_LDG, dscale, params.wgrad_coeff);
+                scaled_write_to_gmem(params.gmem_dbias, thread_c/ELEMENTS_PER_LDG, dbias, params.wgrad_coeff);
+            } else {
+                write_to_gmem(params.gmem_dscale, thread_c/ELEMENTS_PER_LDG, dscale);
+                write_to_gmem(params.gmem_dbias, thread_c/ELEMENTS_PER_LDG, dbias);
+            }
+        }
+
+        // Further normalize the dscale to be used in dx calculation
+        float scale[ELEMENTS_PER_LDG];
+        zero_array(scale);
+        if (is_valid_c) {
+            read_from_gmem(scale, params.gmem_scale, thread_c/ELEMENTS_PER_LDG);
+        }
+        multiply(dscale, var);
+        // scale the inv-var as well, afterwards
+        multiply(var, scale);
+
+        // inverse count
+        float inv_count = params.svar_inv_count;
+
+        // The base pointer to write to.
+        uint16_t *const gmem_dst = &params.gmem_dst[thread_c];
+
+        // Store the elements in registers.
+        #pragma unroll 1
+        for (int loop_i = OUTER_LOOPS-1; loop_i >= 0; --loop_i) {
+            // The value for nhw.
+            int out_nhw = cta_nhw_regs + loop_i*pixels_per_iteration;
+
+            // Normalize the elements and write to memory.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                // Convert to float.
+                float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage[i]);
+                to_float(dy_math, dy_storage[i]);
+                relu_bwd_for_dx(dy_math, x_math, mean_var_scale_bias, var);
+
+                float dx[ELEMENTS_PER_LDG];
+                bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);
+
+                // Write back.
+                const int idx = out_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {
+                    stg_stream(&gmem_dst[idx*params.c], dx);
+                }
+            }
+
+            // The next value of nhw.
+            out_nhw -= pixels_per_iteration;
+
+            // Read the next elements from memory.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                const int idx = out_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {
+                    ldg_stream(x_storage[i], &gmem_src[idx*params.c]);
+                    ldg_stream(dy_storage[i], &gmem_dy[idx*params.c]);
+                }
+            }
+        }
+
+        // Normalize the elements from SMEM and write them out.
+        if (pixels_in_smem > 0) {
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {
+                const int idx = smem_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                const bool is_valid = ((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c;
+                if (is_valid) {
+                    // Read from SMEM.
+                    int offset = i*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                    PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG],
+                        dy_storage_local[PACKED_ELEMENTS_PER_LDG];
+                    read_from_smem(x_storage_local, &smem_storage_packed[offset], threadIdx.x);
+                    offset += PIXELS_PER_THREAD_IN_SMEM*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                    read_from_smem(dy_storage_local, &smem_storage_packed[offset], threadIdx.x);
+                    float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                    to_float(x_math, x_storage_local);
+                    to_float(dy_math, dy_storage_local);
+                    relu_bwd_for_dx(dy_math, x_math, mean_var_scale_bias, var);
+
+                    float dx[ELEMENTS_PER_LDG];
+                    bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);
+
+                    // Write back.
+                    stg_stream(&gmem_dst[idx*params.c], dx);
+                }
+            }
+        }
+        // We're about to start on the next c-blk.  Needed?
+        __syncthreads();
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    typename Storage,
+    int THREADS_PER_CTA,
+    int THREADS_PER_PIXEL,
+    int PIXELS_PER_THREAD_IN_REGISTERS,
+    int PIXELS_PER_THREAD_IN_SMEM,
+    int ELEMENTS_PER_LDG,
+    int USE_ONLINE_APPROACH,
+    int OUTER_LOOPS_,
+    int DESIRED_OCCUPANCY
+>
+__global__ __launch_bounds__(THREADS_PER_CTA, DESIRED_OCCUPANCY)
+    void nhwc_batch_norm_bwd_add_relu(NhwcBatchNormBwdParams params) {
+    // The number of pixels loaded in a single LDG.
+    const int PIXELS_PER_LDG = THREADS_PER_CTA / THREADS_PER_PIXEL;
+    // The number of pixels computed per CTA stored in registers.
+    const int PIXELS_PER_CTA_IN_REGISTERS = PIXELS_PER_THREAD_IN_REGISTERS * PIXELS_PER_LDG;
+    // The number of pixels computed per CTA stored in SMEM.
+    const int PIXELS_PER_CTA_IN_SMEM = PIXELS_PER_THREAD_IN_SMEM*PIXELS_PER_LDG;
+    // The number of C elements per CTA.
+    const int C_ELEMENTS_PER_CTA = THREADS_PER_PIXEL*ELEMENTS_PER_LDG;
+
+    // Shared memory to do CTA-wide parallel sums.
+    __shared__ float smem[THREADS_PER_PIXEL*(THREADS_PER_CTA/32)*ELEMENTS_PER_LDG];
+
+    // The adapter for the storage.
+    typedef PackedStorage<Storage, ELEMENTS_PER_LDG> PackedStorage_;
+    // The data type for packed storage in SMEM.
+    typedef typename PackedStorage_::Type PackedStorageType;
+    // The number of elements in the packed storage.
+    const int PACKED_ELEMENTS_PER_LDG = PackedStorage_::PACKED_ELEMENTS_PER_LDG;
+    // Registers to keep the data live for the persistent approach.
+    PackedStorageType x_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];
+    PackedStorageType dy_storage[PIXELS_PER_THREAD_IN_REGISTERS][PACKED_ELEMENTS_PER_LDG];
+
+    // Shared memory buffer to store the extra pixels.
+    extern __shared__ PackedStorageType smem_storage_packed[];
+
+    for (int c_blk_index = blockIdx.y; c_blk_index < params.c_blks; c_blk_index += gridDim.y) {
+        // The position in the NHW dimension where the CTA starts.
+        int cta_nhw_regs = blockIdx.x * PIXELS_PER_CTA_IN_REGISTERS;
+        // The position in the NHW dimension where the CTA starts for the portion in SMEM.
+        int cta_nhw_smem = blockIdx.x * PIXELS_PER_CTA_IN_SMEM;
+        // Compute the NHW coordinate of the thread in the CTA.
+        const int thread_in_cta_nhw = threadIdx.x / THREADS_PER_PIXEL;
+
+        // The position in the C dimension where the CTA starts.
+        const int cta_c = c_blk_index * C_ELEMENTS_PER_CTA;
+        // Compute the C coordinate of the thread in the CTA.
+        const int thread_in_cta_c = threadIdx.x % THREADS_PER_PIXEL;
+        // Compute the C coordinate of the thread.
+        const int thread_c = cta_c + thread_in_cta_c*ELEMENTS_PER_LDG;
+
+        // Is the thread working on a valid C dimension?
+        const int is_valid_c = thread_c < params.c;
+
+        float mean[ELEMENTS_PER_LDG];
+        zero_array(mean);
+        if (is_valid_c) {
+            read_from_gmem(mean, params.gmem_saved_mean, thread_c/ELEMENTS_PER_LDG);
+        }
+
+        // accumulation related registers
+        float count = 0.f, dscale[ELEMENTS_PER_LDG], dbias[ELEMENTS_PER_LDG];
+        zero_array(dscale);
+        zero_array(dbias);
+
+        // The number of elements loaded by this CTA.
+        int cta_count = 0;
+        // The base pointers to load from.
+        const uint16_t *gmem_src = &params.gmem_src[thread_c];
+        const uint16_t *gmem_dy = &params.gmem_dy[thread_c];
+        uint16_t *gmem_dst1 = &params.gmem_dst1[thread_c];
+
+        // outer loops
+        int OUTER_LOOPS = OUTER_LOOPS_ == 1? 1 : params.outer_loops;
+        // Load the batch of elements. Compute sum across them
+        const int pixels_per_iteration = PIXELS_PER_CTA_IN_REGISTERS*gridDim.x;
+
+        if (OUTER_LOOPS_ != 1) {
+            // We cannot load everything to store persistently, so let's makes sure registers and
+            // smem are fully utilized, offset is evenly divisible by 32
+            int offset = (pixels_per_iteration * OUTER_LOOPS + PIXELS_PER_CTA_IN_SMEM * gridDim.x -
+                          params.nhw) & ~31;
+            cta_nhw_regs -= offset;
+            cta_nhw_smem -= offset;
+        }
+
+        const unsigned int *const gmem_relu_bitmask = params.gmem_relu_bitmask +
+                                      ((params.nhw + 31) & ~31) * 2 * c_blk_index;
+
+        #pragma unroll 1
+        for (int loop_i = 0; loop_i < OUTER_LOOPS; ++loop_i) {
+            // The nhw position.
+            int nhw_regs = cta_nhw_regs + loop_i*pixels_per_iteration;
+            // Update the number of elements loaded by this CTA. TODO: Skip if <= 0!!!
+            cta_count += max(0, min(PIXELS_PER_CTA_IN_REGISTERS, params.nhw-nhw_regs));
+
+            int lane_id = threadIdx.x & 31;
+
+            // Read the elements from memory.
+            float is_valid[PIXELS_PER_THREAD_IN_REGISTERS];
+            unsigned int relu_mask[PIXELS_PER_THREAD_IN_REGISTERS];
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                const int idx = nhw_regs + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                zero_array(x_storage[i]);
+                zero_array(dy_storage[i]);
+                is_valid[i] = 0.f;
+                const bool is_valid_nhw =
+                    static_cast<unsigned int>(idx) < static_cast<unsigned int>(params.nhw);
+                if (is_valid_nhw) {
+                    if (is_valid_c) {
+                        if (loop_i == OUTER_LOOPS - 1) {
+                            ldg_stream(x_storage[i], &gmem_src[idx*params.c]);
+                            ldg_stream(dy_storage[i], &gmem_dy[idx*params.c]);
+                        } else {
+                            ldg(x_storage[i], &gmem_src[idx*params.c]);
+                            ldg(dy_storage[i], &gmem_dy[idx*params.c]);
+                        }
+                        is_valid[i] = 1.f;
+                    }
+
+                    if (lane_id < ELEMENTS_PER_LDG) {
+                        relu_mask[i] = gmem_relu_bitmask[idx * 2 + lane_id];
+                    }
+                }
+            }
+
+            // Do the math.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                const int idx = nhw_regs + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                // Convert to float and update
+                float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                bool rectified[ELEMENTS_PER_LDG];
+                #pragma unroll
+                for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {
+                    rectified[j] = ((__shfl_sync(0xFFFFFFFFU, relu_mask[i], j) &
+                                    (1U << lane_id)) != 0);
+                }
+                to_float(x_math, x_storage[i]);
+                to_float(dy_math, dy_storage[i]);
+
+                // Update the count.
+                count += is_valid[i];
+                // Invert the count.
+                float inv_count = is_valid[i] ? 1.f / count : 0.f;
+
+                relu_bwd(dy_math, rectified, is_valid[i]);
+                bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);
+
+                // Lastly we need 'dy' only for BN, so store the 'relu-dgrad'ed version
+                from_float(dy_storage[i], dy_math);
+
+                // dZ for elementwise add
+                if (is_valid[i]) {
+                    if (loop_i == OUTER_LOOPS - 1) {
+                        stg_stream(&gmem_dst1[idx*params.c], dy_storage[i]);
+                    } else {
+                        stg(&gmem_dst1[idx*params.c], dy_storage[i]);
+                    }
+                }
+            }
+        }
+
+        // The elements to load and store in SMEM.
+        int smem_nhw = OUTER_LOOPS*pixels_per_iteration + cta_nhw_smem;
+        // Load elements from SMEM, update the CTA count.
+        int pixels_in_smem = min(PIXELS_PER_CTA_IN_SMEM, params.nhw-smem_nhw);
+        if (pixels_in_smem > 0) {
+            cta_count += pixels_in_smem;
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {
+                const int idx = smem_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                const bool is_pixel_valid_nhw =
+                    static_cast<unsigned int>(idx) < static_cast<unsigned int>(params.nhw);
+                const bool is_pixel_valid = is_pixel_valid_nhw && is_valid_c;
+                PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG],
+                                  dy_storage_local[PACKED_ELEMENTS_PER_LDG];
+                unsigned int relu_mask;
+                int lane_id = threadIdx.x & 31;
+                zero_array(x_storage_local);
+                zero_array(dy_storage_local);
+                if (is_pixel_valid_nhw) {
+                    if (is_valid_c) {
+                        ldg_stream(x_storage_local, &gmem_src[idx*params.c]);
+                        ldg_stream(dy_storage_local, &gmem_dy[idx*params.c]);
+                    }
+                    if (lane_id < ELEMENTS_PER_LDG) {
+                        relu_mask = gmem_relu_bitmask[idx * 2 + lane_id];
+                    }
+                }
+                bool rectified[ELEMENTS_PER_LDG];
+                #pragma unroll
+                for (int j = 0; j < ELEMENTS_PER_LDG; ++j) {
+                    rectified[j] = ((__shfl_sync(0xFFFFFFFFU, relu_mask, j) &
+                                    (1U << lane_id)) != 0);
+                }
+
+                // The offset to store in SMEM.
+                int offset = i*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                // Store in SMEM.
+                write_to_smem(&smem_storage_packed[offset], threadIdx.x, x_storage_local);
+                offset += PIXELS_PER_THREAD_IN_SMEM*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                // Update the count.
+                count += is_pixel_valid;
+                // Invert the count.
+                float inv_count = is_pixel_valid ? 1.f / count : 0.f;
+
+                float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage_local);
+                to_float(dy_math, dy_storage_local);
+
+                relu_bwd(dy_math, rectified, is_pixel_valid);
+                bwd_update(dscale, dbias, dy_math, x_math, mean, inv_count);
+
+                from_float(dy_storage_local, dy_math);
+                // dZ for elementwise add
+                if (is_pixel_valid) {
+                    stg_stream(&gmem_dst1[idx*params.c], dy_storage_local);
+                }
+                // only store the 'relu-dgrad'ed version!
+                write_to_smem(&smem_storage_packed[offset], threadIdx.x, dy_storage_local);
+            }
+        }
+
+        // We scale the mean by the number of elements. It brings more stability.
+        #pragma unroll
+        for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+            dbias[i] *= count;
+            dscale[i] *= count;
+        }
+
+        // dscale parallel sum
+        ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+            smem, dscale, thread_in_cta_nhw);
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dscale, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // dbias parallel sum
+        ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+            smem, dbias, thread_in_cta_nhw);
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dbias, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // The workspace in global memory is distributed across the different CTA.
+        int gmem_sums_offset = c_blk_index*gridDim.x*C_ELEMENTS_PER_CTA*2;
+        // Write the data for the CTA to global memory.
+        float *gmem_sums = &params.gmem_sums[gmem_sums_offset];
+        if (threadIdx.x < THREADS_PER_PIXEL) {
+            const int idx = blockIdx.x*THREADS_PER_PIXEL + threadIdx.x;
+            write_to_gmem(&gmem_sums[                           0], idx, dscale);
+            write_to_gmem(&gmem_sums[C_ELEMENTS_PER_CTA*gridDim.x], idx, dbias);
+        }
+
+        // The counters to count how many CTAs have retired at this point.
+        // A given cta uses the same counter every other time through the outer loop.
+        int *gmem_retired_ctas = &params.gmem_retired_ctas[c_blk_index % (2 * gridDim.y)];
+        inter_block_sync(gmem_retired_ctas, gridDim.x, blockIdx.x == 0);
+
+        // Reset the accumulators for global summation
+        zero_array(dscale);
+        zero_array(dbias);
+
+        // Build the global accumulation
+        #pragma unroll 1
+        for (int idx = threadIdx.x; idx < THREADS_PER_PIXEL*gridDim.x; idx += THREADS_PER_CTA) {
+            float tmp1[ELEMENTS_PER_LDG], tmp2[ELEMENTS_PER_LDG];
+            read_from_gmem(tmp1, gmem_sums,                              idx);
+            read_from_gmem(tmp2, gmem_sums+C_ELEMENTS_PER_CTA*gridDim.x, idx);
+
+            #pragma unroll
+            for (int i = 0; i < ELEMENTS_PER_LDG; ++i) {
+                dscale[i] += tmp1[i];
+                dbias[i] += tmp2[i];
+            }
+        }
+
+        // dscale parallel sum
+        if (params.sync_iters>0) {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatchX<THREADS_PER_CTA>(
+                smem, dscale, thread_in_cta_nhw, params.my_data, params.pair_datas, 4*c_blk_index+1, params.magic, params.sync_iters);
+        } else {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+                smem, dscale, thread_in_cta_nhw);
+        }
+
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dscale, smem, thread_in_cta_c);
+        __syncthreads();
+
+        // dbias parallel sum
+        if (params.sync_iters>0) {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatchX<THREADS_PER_CTA>(
+                smem, dbias, thread_in_cta_nhw, params.my_data, params.pair_datas, 4*c_blk_index+0, params.magic, params.sync_iters);
+        } else {
+            ParallelSums<THREADS_PER_PIXEL, ELEMENTS_PER_LDG>::dispatch<THREADS_PER_CTA>(
+                smem, dbias, thread_in_cta_nhw);
+        }
+
+        __syncthreads();
+        // The values in shared memory correspond to the CTA-wide sums.
+        read_from_smem(dbias, smem, thread_in_cta_c);
+
+        // Normalize the dscale.
+        float var[ELEMENTS_PER_LDG];
+        zero_array(var);
+        if (is_valid_c) {
+            read_from_gmem(var, params.gmem_saved_var, thread_c/ELEMENTS_PER_LDG);
+        }
+        multiply(dscale, var);
+
+        // store dscale/dbias
+        bool is_valid_for_saving = is_valid_c && blockIdx.x == 0 && thread_in_cta_nhw == 0;
+        if (is_valid_for_saving) {
+            if (params.sync_iters>0)
+            {
+                scaled_write_to_gmem(params.gmem_dscale, thread_c/ELEMENTS_PER_LDG, dscale, params.wgrad_coeff);
+                scaled_write_to_gmem(params.gmem_dbias, thread_c/ELEMENTS_PER_LDG, dbias, params.wgrad_coeff);
+            } else {
+                write_to_gmem(params.gmem_dscale, thread_c/ELEMENTS_PER_LDG, dscale);
+                write_to_gmem(params.gmem_dbias, thread_c/ELEMENTS_PER_LDG, dbias);
+            }
+        }
+
+        // Further normalize the dscale to be used in dx calculation
+        float scale[ELEMENTS_PER_LDG];
+        zero_array(scale);
+        if (is_valid_c) {
+            read_from_gmem(scale, params.gmem_scale, thread_c/ELEMENTS_PER_LDG);
+        }
+        multiply(dscale, var);
+        // scale the inv-var as well, afterwards
+        multiply(var, scale);
+
+        // inverse count
+        float inv_count = params.svar_inv_count;
+
+        // The base pointer to write to.
+        uint16_t *const gmem_dst = &params.gmem_dst[thread_c];
+
+        // Store the elements in registers.
+        #pragma unroll 1
+        for (int loop_i = OUTER_LOOPS-1; loop_i >= 0; --loop_i) {
+            // The value for nhw.
+            int out_nhw = cta_nhw_regs + loop_i*pixels_per_iteration;
+
+            // Normalize the elements and write to memory.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                const int idx = out_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                const bool is_valid = ((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c;
+                // Convert to float.
+                float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                to_float(x_math, x_storage[i]);
+                to_float(dy_math, dy_storage[i]);
+
+                float dx[ELEMENTS_PER_LDG];
+                bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);
+
+                // Write back.
+                if (is_valid) {
+                    stg_stream(&gmem_dst[idx*params.c], dx);
+                }
+            }
+
+            // The next value of nhw.
+            out_nhw -= pixels_per_iteration;
+
+            // Read the next elements from memory.
+            #pragma unroll
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_REGISTERS; ++i) {
+                const int idx = out_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                float y[ELEMENTS_PER_LDG];
+                zero_array(y);
+                if (((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c) {
+                    ldg_stream(x_storage[i], &gmem_src[idx*params.c]);
+                    ldg_stream(dy_storage[i], &gmem_dst1[idx*params.c]);
+                }
+            }
+        }
+
+        // Normalize the elements from SMEM and write them out.
+        if (pixels_in_smem > 0) {
+            for (int i = 0; i < PIXELS_PER_THREAD_IN_SMEM; ++i) {
+                const int idx = smem_nhw + thread_in_cta_nhw + i*PIXELS_PER_LDG;
+                const bool is_valid = ((unsigned int)idx < (unsigned int)params.nhw) && is_valid_c;
+                if (is_valid) {
+                    // Read from SMEM.
+                    int offset = i*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                    PackedStorageType x_storage_local[PACKED_ELEMENTS_PER_LDG],
+                        dy_storage_local[PACKED_ELEMENTS_PER_LDG];
+                    read_from_smem(x_storage_local, &smem_storage_packed[offset], threadIdx.x);
+                    offset += PIXELS_PER_THREAD_IN_SMEM*THREADS_PER_CTA*PACKED_ELEMENTS_PER_LDG;
+                    read_from_smem(dy_storage_local, &smem_storage_packed[offset], threadIdx.x);
+                    float x_math[ELEMENTS_PER_LDG], dy_math[ELEMENTS_PER_LDG];
+                    to_float(x_math, x_storage_local);
+                    to_float(dy_math, dy_storage_local);
+
+                    float dx[ELEMENTS_PER_LDG];
+                    bwd_dx(dx, dy_math, var, x_math, mean, dscale, dbias, inv_count);
+
+                    // Write back.
+                    stg_stream(&gmem_dst[idx*params.c], dx);
+                }
+            }
+        }
+        // We're about to start on the next c-blk.  Needed?
+        __syncthreads();
+    }
+}
+
+#endif  // MXNET_OPERATOR_NN_CUDNN_NHWC_BATCH_NORM_KERNEL_H_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln.h
new file mode 100644
index 0000000000..07392a192b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln.h
@@ -0,0 +1,200 @@
+#pragma once
+
+#include <unordered_map>
+#include <cuda_fp16.h>
+#include <cuda_bf16.h>
+
+namespace layer_norm {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Params> 
+struct LaunchParams{
+
+    size_t workspace_bytes;
+    size_t barrier_size;
+
+    cudaDeviceProp * props;
+
+    cudaStream_t stream;
+
+    Params params;
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct ParamsBase {
+    ParamsBase()
+        : ctas_per_col(0)
+        , rows(0)
+        , cols(0)
+        , x(nullptr)
+        , mu(nullptr)
+        , rs(nullptr)
+        , gamma(nullptr)
+        , workspace(nullptr)
+        , barrier(nullptr)
+    {
+    }
+
+    // For Multi-CTA, number of different CTA groups. Otherwise same as gridDim.x.
+    int ctas_per_col;
+
+    // Input is interpreted as matrix. We normalize across columns.
+    int rows;
+    int cols;
+
+    // Common data pointers.
+    void *x;
+    void *mu;
+    void *rs;
+    void *gamma;
+
+    // Multi-CTA workspace in gmem.
+    void *workspace;
+
+    // Multi-CTA sync barriers in gmem.
+    int *barrier;
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct FwdParams : public ParamsBase {
+    FwdParams()
+        : ParamsBase()
+        , z(nullptr)
+        , beta(nullptr)
+        , epsilon(0.f)
+    {
+    }
+
+    // Output of LN FWD.
+    void *z;
+    void *beta;
+    float epsilon;
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct BwdParams : public ParamsBase {
+    BwdParams()
+        : ParamsBase()
+        , dz(nullptr)
+        , dbeta_part(nullptr)
+        , dgamma_part(nullptr)
+        , dx(nullptr)
+        , dbeta(nullptr)
+        , dgamma(nullptr)
+    {
+    }
+
+    // Input: gradient wrt. LN FWD output.
+    void *dz;
+
+    // Workspace for Wgrad pre-reduction.
+    void *dbeta_part;
+    void *dgamma_part;
+
+    // Output: Dgrad.
+    void *dx;
+    // Output: Wgrad.
+    void *dbeta;
+    void *dgamma;
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+using FwdFunction = std::function<void(LaunchParams<FwdParams>&, const bool)>;
+using BwdFunction = std::function<void(LaunchParams<BwdParams>&, const bool)>;
+using FunctionKey = uint64_t;
+using FwdRegistry = std::unordered_map<FunctionKey, FwdFunction>;
+using BwdRegistry = std::unordered_map<FunctionKey, BwdFunction>;
+
+extern FwdRegistry FWD_FUNCS;
+extern BwdRegistry BWD_FUNCS;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+using fp32 = float;
+using fp16 = half;
+using bf16 = nv_bfloat16;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+struct TypeId{};
+
+template<>
+struct TypeId<fp16>{
+    constexpr static uint32_t Value = 0;
+};
+
+template<>
+struct TypeId<bf16>{
+    constexpr static uint32_t Value = 1;
+};
+
+template<>
+struct TypeId<fp32>{
+    constexpr static uint32_t Value = 2;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, int S>
+struct Type2Key{
+    constexpr static uint32_t Value = TypeId<T>::Value << S;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+struct WeightType2Key : public Type2Key<T, 0>{};
+
+template<typename T>
+struct InputType2Key : public Type2Key<T, 2>{};
+
+template<typename T>
+struct OutputType2Key : public Type2Key<T, 4>{};
+
+template<typename T>
+struct ComputeType2Key : public Type2Key<T, 6>{};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename W, typename I, typename O, typename C>
+struct Types2Key{
+    constexpr static uint32_t Value = WeightType2Key<W>::Value | InputType2Key<I>::Value | OutputType2Key<O>::Value | ComputeType2Key<C>::Value;
+    constexpr static inline uint64_t get(const uint64_t hidden_size){
+        constexpr uint64_t type_key = Value;
+        return (type_key << 32) | hidden_size;
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename W, typename I, typename O, typename C, uint64_t HIDDEN_SIZE>
+struct FwdRegistrar{
+    FwdRegistrar(FwdFunction f){
+        uint64_t key = Types2Key<W,I,O,C>::get(HIDDEN_SIZE);
+        FWD_FUNCS.insert({ key, f });
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename W, typename I, typename O, typename C, uint64_t HIDDEN_SIZE>
+struct BwdRegistrar{
+    BwdRegistrar(BwdFunction f){
+        uint64_t key = Types2Key<W,I,O,C>::get(HIDDEN_SIZE);
+        BWD_FUNCS.insert({ key, f });
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace layer_norm
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_api.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_api.cpp
new file mode 100644
index 0000000000..30e4a5fec6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_api.cpp
@@ -0,0 +1,246 @@
+#include <torch/extension.h>
+#include "ATen/cuda/CUDAContext.h"
+
+#include "ln.h"
+
+/*
+
+Supported Type combinations:
+
+input    compute   weights   output    
+=======================================
+fp32     fp32      fp32      fp32      
+fp16     fp32      fp16      fp16      
+bf16     fp32      bf16      bf16      
+fp32     fp32      fp16      fp16      
+fp32     fp32      bf16      bf16      
+
+Remarks:
+Output type = Weight type
+Compute always in FP32
+
+*/
+
+namespace layer_norm {
+
+// Create registries and provide runtime versions of config hash functions.
+
+FwdRegistry FWD_FUNCS;
+BwdRegistry BWD_FUNCS;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+uint32_t get_type_id(torch::Dtype dtype){
+    if( dtype == torch::kFloat16 ) {
+        return TypeId<fp16>::Value;
+    } else if( dtype == torch::kBFloat16 ) {
+        return TypeId<bf16>::Value;
+    } else if( dtype == torch::kFloat32 ) {
+        return TypeId<fp32>::Value;
+    } else {
+        TORCH_CHECK(false, "Type not supported: ", dtype);
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+uint64_t get_key(torch::Dtype wtype, torch::Dtype itype, torch::Dtype otype, torch::Dtype ctype, uint64_t hidden_size) {
+    using namespace layer_norm;
+    uint64_t type_key = get_type_id(wtype) | (get_type_id(itype) << 2) | (get_type_id(otype) << 4) | (get_type_id(ctype) << 6);
+    uint64_t launcher_key = (type_key << 32) | hidden_size;
+    return launcher_key;
+}
+
+}  // namespace layer_norm
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+layer_norm::FwdFunction & get_fwd_launcher(torch::Dtype wtype, torch::Dtype itype, torch::Dtype otype, torch::Dtype ctype, uint32_t hidden_size) {
+    auto iter = layer_norm::FWD_FUNCS.find(layer_norm::get_key(wtype, itype, otype, ctype, hidden_size));
+    if( iter != layer_norm::FWD_FUNCS.end() ) {
+        return iter->second;
+    } else {
+        TORCH_CHECK(false, "FWD: Unsupported hidden_size or types: ", hidden_size, wtype, itype, otype, ctype);
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+layer_norm::BwdFunction & get_bwd_launcher(torch::Dtype wtype, torch::Dtype itype, torch::Dtype otype, torch::Dtype ctype, uint32_t hidden_size) {
+    auto iter = layer_norm::BWD_FUNCS.find(layer_norm::get_key(wtype, itype, otype, ctype, hidden_size));
+    if( iter != layer_norm::BWD_FUNCS.end() ) {
+        return iter->second;
+    } else {
+        TORCH_CHECK(false, "BWD: Unsupported hidden_size or types: ", hidden_size, wtype, itype, otype, ctype);
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+std::vector<at::Tensor> ln_fwd(const at::Tensor &x,      // BxSxhidden_size
+                               const at::Tensor &gamma,   // hidden_size
+                               const at::Tensor &beta,   // hidden_size
+                               const float epsilon
+) {
+    auto itype = x.scalar_type();
+    auto wtype = gamma.scalar_type();
+    auto otype = wtype;
+    auto ctype = torch::kFloat32;
+
+    TORCH_CHECK(beta.scalar_type() == wtype);
+
+    TORCH_CHECK(x.is_cuda())
+    TORCH_CHECK(gamma.is_cuda())
+    TORCH_CHECK(beta.is_cuda())
+
+    TORCH_CHECK(x.is_contiguous());
+    auto sizes = x.sizes();
+    TORCH_CHECK(sizes.size() == 2);
+
+    const int rows = sizes[0];
+    const int cols = sizes[1];
+    auto hidden_size = gamma.numel();
+
+    TORCH_CHECK(gamma.sizes() == beta.sizes());
+    TORCH_CHECK(hidden_size == cols);
+
+    TORCH_CHECK(epsilon >= 0.f);
+
+    auto opts = x.options();
+
+    auto z = torch::empty(sizes, opts.dtype(otype));
+
+    auto mu = torch::empty({ rows }, opts.dtype(ctype));
+    auto rsigma = torch::empty({ rows }, opts.dtype(ctype));
+
+    layer_norm::LaunchParams<layer_norm::FwdParams> launch_params;
+
+    launch_params.props = at::cuda::getCurrentDeviceProperties();
+    launch_params.stream = at::cuda::getCurrentCUDAStream().stream();
+
+    // Request the kernel launcher.
+    auto launcher = get_fwd_launcher(wtype, itype, otype, ctype, hidden_size);
+
+    // Query the kernel-specific launch parameters.
+    launcher(launch_params, true);
+
+    at::Tensor workspace, barrier;
+
+    // Set the kernel runtime parameters.
+    layer_norm::FwdParams &params = launch_params.params;
+    params.rows = rows;
+    params.cols = cols;
+    params.x = x.data_ptr();
+    params.mu = mu.data_ptr();
+    params.rs = rsigma.data_ptr();
+    params.gamma = gamma.data_ptr();
+    params.beta = beta.data_ptr();
+    params.z = z.data_ptr();
+    params.epsilon = epsilon;
+
+    if( launch_params.barrier_size > 0 ) {
+        auto options = x.options();
+        barrier = torch::zeros(launch_params.barrier_size, options.dtype(torch::kInt32));
+        workspace = torch::empty(launch_params.workspace_bytes, options.dtype(torch::kChar));
+        params.workspace = workspace.data_ptr();
+        params.barrier = barrier.data_ptr<int>();
+    }
+
+    // Launch the kernel.
+    launcher(launch_params, false);
+
+    return { z, mu, rsigma };
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+std::vector<at::Tensor> ln_bwd(const at::Tensor &dz,     // BxSxhidden_size
+                               const at::Tensor &x,      // BxSxhidden_size
+                               const at::Tensor &mu,     // BxS, FP32!
+                               const at::Tensor &rsigma, // BxS, FP32!
+                               const at::Tensor &gamma   // hidden_size
+) {
+
+    auto itype = x.scalar_type();
+    auto wtype = gamma.scalar_type();
+    auto otype = wtype;
+    auto ctype = torch::kFloat32;
+
+    TORCH_CHECK(dz.dtype() == otype);
+    TORCH_CHECK(mu.dtype() == ctype);
+    TORCH_CHECK(rsigma.dtype() == ctype);
+
+    TORCH_CHECK(x.is_cuda());
+    TORCH_CHECK(dz.is_cuda());
+    TORCH_CHECK(mu.is_cuda());
+    TORCH_CHECK(rsigma.is_cuda());
+    TORCH_CHECK(gamma.is_cuda());
+
+    TORCH_CHECK(x.is_contiguous());
+    TORCH_CHECK(dz.is_contiguous());
+
+    auto sizes = x.sizes();
+    TORCH_CHECK(sizes.size() == 2);
+    TORCH_CHECK(dz.sizes() == sizes);
+    auto rows = sizes[0];
+    auto cols = sizes[1];
+
+    auto hidden_size = gamma.numel();
+
+    TORCH_CHECK(mu.numel() == rows);
+    TORCH_CHECK(mu.sizes() == rsigma.sizes());
+
+    TORCH_CHECK(gamma.numel() == cols);
+
+    auto options = x.options();
+
+    auto dx = torch::empty_like(x);
+    auto dgamma = torch::empty_like(gamma);
+    auto dbeta = torch::empty_like(gamma);
+
+    layer_norm::LaunchParams<layer_norm::BwdParams> launch_params;
+    launch_params.stream = at::cuda::getCurrentCUDAStream().stream();
+    launch_params.props = at::cuda::getCurrentDeviceProperties();
+
+    auto launcher = get_bwd_launcher(wtype, itype, otype, ctype, hidden_size);
+
+    launcher(launch_params, true);
+
+    auto dgamma_part = torch::empty({ launch_params.params.ctas_per_col, hidden_size }, options.dtype(ctype));
+    auto dbeta_part = torch::empty({ launch_params.params.ctas_per_col, hidden_size }, options.dtype(ctype));
+    at::Tensor workspace, barrier;
+
+    layer_norm::BwdParams &params = launch_params.params;
+    params.rows = rows;
+    params.cols = cols;
+    params.x = x.data_ptr();
+    params.mu = mu.data_ptr();
+    params.rs = rsigma.data_ptr();
+    params.gamma = gamma.data_ptr();
+    params.dz = dz.data_ptr();
+    params.dx = dx.data_ptr();
+    params.dbeta = dbeta.data_ptr();
+    params.dgamma = dgamma.data_ptr();
+    params.dbeta_part = dbeta_part.data_ptr();
+    params.dgamma_part = dgamma_part.data_ptr();
+
+    if( launch_params.barrier_size > 0 ) {
+        // TODO Any way to avoid this?
+        barrier = torch::zeros(launch_params.barrier_size, options.dtype(torch::kInt32));
+        workspace = torch::empty(launch_params.workspace_bytes, options.dtype(torch::kChar));
+        params.workspace = workspace.data_ptr();
+        params.barrier = barrier.data_ptr<int>();
+    }
+
+    launcher(launch_params, false);
+
+    return { dx, dgamma, dbeta, dgamma_part, dbeta_part };
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.doc() = "CUDA LayerNorm"; 
+  m.def("ln_fwd", &ln_fwd, "Run LayerNorm forward kernel");
+  m.def("ln_bwd", &ln_bwd, "Run LayerNorm backward kernel");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_bwd_kernels.cuh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_bwd_kernels.cuh
new file mode 100644
index 0000000000..8595f5ed44
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_bwd_kernels.cuh
@@ -0,0 +1,315 @@
+#pragma once
+
+namespace layer_norm {
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) 
+void ln_bwd_kernel(layer_norm::BwdParams params) {
+
+    enum { ROWS_PER_CTA = Ktraits::ROWS_PER_CTA };
+    enum { WARPS_M = Ktraits::WARPS_M };
+    enum { WARPS_N = Ktraits::WARPS_N };
+    enum { THREADS_PER_ROW = Ktraits::THREADS_PER_ROW };
+    enum { COLS = Ktraits::COLS };
+    enum { BYTES_PER_ROW = Ktraits::BYTES_PER_ROW };
+    enum { LDGS = Ktraits::LDGS };
+    enum { NUM_ELTS = Ktraits::ELTS_PER_LDG };
+    enum { THREADS_PER_WARP = Ktraits::THREADS_PER_WARP };
+    enum { CTAS_PER_ROW = Ktraits::CTAS_PER_ROW };
+
+    using compute_t = typename Ktraits::compute_t;
+    using index_t = typename Ktraits::index_t;
+    using Ivec = typename Ktraits::Ivec;
+    using Ovec = typename Ktraits::Ovec;
+    using Wvec = typename Ktraits::Wvec;
+    using Cvec = typename Ktraits::Cvec;
+    using Reducer = typename Ktraits::Reducer;
+    using reduce_t = typename Reducer::Type;
+
+    extern __shared__ char smem_[];
+
+    const index_t tidx = threadIdx.x;
+    const index_t bidn = blockIdx.x % CTAS_PER_ROW;
+    const index_t bidm = blockIdx.x / CTAS_PER_ROW;
+    const index_t lane = tidx % THREADS_PER_WARP;
+    const index_t warp = tidx / THREADS_PER_WARP;
+    const index_t warp_m = warp / Ktraits::WARPS_N;
+    const index_t warp_n = warp % Ktraits::WARPS_N;
+    const index_t tid_r = warp_n * THREADS_PER_WARP + lane;
+
+    const index_t r = bidm * Ktraits::ROWS_PER_CTA + warp_m;
+    const index_t c = bidn * THREADS_PER_ROW + warp_n * THREADS_PER_WARP + lane;
+
+    static_assert(COLS == THREADS_PER_ROW * LDGS * NUM_ELTS * CTAS_PER_ROW);
+
+    Cvec dzy_sum[LDGS];
+    Cvec dz_sum[LDGS];
+
+    memset(dzy_sum, 0, sizeof(dzy_sum));
+    memset(dz_sum, 0, sizeof(dz_sum));
+
+    compute_t * smem_wgrad = reinterpret_cast<compute_t*>(smem_);
+    char *smem_dgrad = smem_ + Ktraits::SMEM_BYTES_WGRAD;
+
+    Reducer reducer(params, bidm, bidn, warp_m, warp_n, lane, smem_dgrad);
+
+    Sum<reduce_t> sum;
+
+    constexpr float rn = 1.f / float(COLS);
+    Wvec gamma[LDGS];
+    index_t idx = c;
+    #pragma unroll
+    for( int it = 0; it < LDGS; it++ ) {
+        gamma[it].load_from(params.gamma, idx);
+        idx += Ktraits::VEC_COLS_PER_LDG;
+    }
+    // TODO if ROWS_PER_CTA does not divide rows, we might get divergence in the
+    // last blocks with syncthreads!
+    // grid stride over rows
+    #pragma unroll 1
+    for( int row = r; row < params.rows; row += params.ctas_per_col * ROWS_PER_CTA ) {
+        const compute_t mu_r = static_cast<const compute_t *>(params.mu)[row];
+        const compute_t rs_r = static_cast<const compute_t *>(params.rs)[row];
+        Ivec x[LDGS];
+        Ovec dz[LDGS];
+        index_t idx = row * Ktraits::VEC_COLS + c;
+        #pragma unroll
+        for( int it = 0; it < LDGS; it++ ) {
+            dz[it].load_from(params.dz, idx);
+            x[it].load_from(params.x, idx);
+            idx += Ktraits::VEC_COLS_PER_LDG;
+        }
+
+        compute_t dy[LDGS * NUM_ELTS];
+        compute_t y[LDGS * NUM_ELTS];
+
+        compute_t mdy_local = 0.f;
+        compute_t mdyy_local = 0.f;
+        #pragma unroll
+        for( int it = 0; it < LDGS; it++ ) {
+            #pragma unroll
+            for( int jt = 0; jt < NUM_ELTS; jt++ ) {
+                compute_t x_tmp = x[it].data.elt[jt];
+                compute_t y_tmp = rs_r * (x_tmp - mu_r);
+                compute_t dy_tmp = compute_t(gamma[it].data.elt[jt]);
+                dy_tmp *= compute_t(dz[it].data.elt[jt]);
+                compute_t dz_tmp = dz[it].data.elt[jt];
+
+                mdy_local += dy_tmp;
+                mdyy_local += dy_tmp * y_tmp;
+
+                dy[it * NUM_ELTS + jt] = dy_tmp;
+                y[it * NUM_ELTS + jt] = y_tmp;
+
+                dzy_sum[it].data.elt[jt] += dz_tmp * y_tmp;
+                dz_sum[it].data.elt[jt] += dz_tmp;
+            }
+        }
+
+        reduce_t result = reducer.allreduce({mdy_local, mdyy_local}, sum);
+        mdy_local = layer_norm::Get<0>::of<reduce_t, compute_t>(result) * rn;
+        mdyy_local = layer_norm::Get<1>::of<reduce_t, compute_t>(result) * rn;
+
+        Ivec dx[LDGS];
+        idx = row * Ktraits::VEC_COLS + c;
+        #pragma unroll
+        for( int it = 0; it < LDGS; it++ ) {
+            #pragma unroll
+            for( int jt = 0; jt < NUM_ELTS; jt++ ) {
+                compute_t dy_tmp = dy[it * NUM_ELTS + jt];
+                compute_t y_tmp = y[it * NUM_ELTS + jt];
+                compute_t dx_tmp = rs_r * (dy_tmp - (mdyy_local * y_tmp + mdy_local));
+                dx[it].data.elt[jt] = dx_tmp;
+            }
+            dx[it].store_to(params.dx, idx);
+            idx += Ktraits::VEC_COLS_PER_LDG;
+        }
+
+    }  // end: grid stride loop
+
+    if( WARPS_M == 1 ) {
+        idx = r * Ktraits::VEC_COLS + c;
+        #pragma unroll
+        for( int it = 0; it < LDGS; it++ ) {
+            dz_sum[it].store_to(params.dbeta_part, idx);
+            dzy_sum[it].store_to(params.dgamma_part, idx);
+            idx += Ktraits::VEC_COLS_PER_LDG;
+        }
+    } else {
+        static_assert(WARPS_M == 1 || Ktraits::CTAS_PER_ROW == 1, "Multiple rows per CTA not supported for Multi-CTA.");
+        // Finalize reduction of part dgamma and dbeta for this CTA
+        // by reducing over the rows held across the WARPS_M warps
+
+        // Assumption: blockSize divides hidden size.
+        enum { NUM_RES = COLS / Ktraits::THREADS_PER_CTA };
+        static_assert(NUM_RES * Ktraits::THREADS_PER_CTA == COLS, "");
+
+        idx = warp_m * Ktraits::VEC_COLS + tid_r;
+        #pragma unroll
+        for( int it = 0; it < LDGS; it++ ) {
+            dz_sum[it].store_to(smem_wgrad, idx);
+            idx += THREADS_PER_ROW;
+        }
+        __syncthreads();
+        compute_t cta_dz_sum[NUM_RES];
+        memset(cta_dz_sum, 0, sizeof(compute_t) * NUM_RES);
+        for( int it = 0; it < ROWS_PER_CTA; it++ ) {
+            for( int jt = 0; jt < NUM_RES; jt++ ) {
+                cta_dz_sum[jt] += smem_wgrad[it * COLS + tidx + jt * Ktraits::THREADS_PER_CTA];
+            }
+        }
+        __syncthreads();
+
+        idx = warp_m * Ktraits::VEC_COLS + tid_r;
+        #pragma unroll
+        for( int it = 0; it < LDGS; it++ ) {
+            dzy_sum[it].store_to(smem_wgrad, idx);
+            idx += THREADS_PER_ROW;
+        }
+        __syncthreads();
+        compute_t cta_dzy_sum[NUM_RES];
+        memset(cta_dzy_sum, 0, sizeof(compute_t) * NUM_RES);
+        for( int it = 0; it < ROWS_PER_CTA; it++ ) {
+            for( int jt = 0; jt < NUM_RES; jt++ ) {
+                cta_dzy_sum[jt] += smem_wgrad[it * COLS + tidx + jt * Ktraits::THREADS_PER_CTA];
+            }
+        }
+
+        compute_t *dgamma_part = static_cast<compute_t *>(params.dgamma_part) + bidm * COLS + tidx;
+        for( int jt = 0; jt < NUM_RES; jt++ ) {
+            *dgamma_part = cta_dzy_sum[jt];
+            dgamma_part += Ktraits::THREADS_PER_CTA;
+        }
+
+        compute_t *dbeta_part = static_cast<compute_t *>(params.dbeta_part) + bidm * COLS + tidx;
+        for( int jt = 0; jt < NUM_RES; jt++ ) {
+            *dbeta_part = cta_dz_sum[jt];
+            dbeta_part += Ktraits::THREADS_PER_CTA;
+        }
+    }
+}
+
+template<typename Kernel_traits>
+__global__ __launch_bounds__(Kernel_traits::THREADS_PER_CTA)
+void ln_bwd_finalize_kernel(BwdParams params)
+{
+
+    using compute_t = typename Kernel_traits::compute_t;
+    using weight_t = typename Kernel_traits::weight_t;
+    using index_t = typename Kernel_traits::index_t;
+    using Reducer = typename Kernel_traits::Reducer;
+    using reduce_t = typename Reducer::Type;
+
+    Sum<reduce_t> sum;
+    enum { NUM_ELT = Kernel_traits::ELTS_PER_LDG };
+    enum { THREADS_PER_WARP = Kernel_traits::THREADS_PER_WARP };
+
+    __shared__ char smem_[Kernel_traits::SMEM_BYTES_PER_CTA];
+
+    constexpr uint32_t bidm = 0;
+
+    const uint32_t bidn = blockIdx.x;
+    const uint32_t tidx = threadIdx.x;
+    const uint32_t warp = tidx / THREADS_PER_WARP;
+    const uint32_t lane = tidx % THREADS_PER_WARP;
+
+    Reducer reducer(params, bidm, bidn, 0, 0, lane, smem_);
+
+    const uint32_t c = bidn * THREADS_PER_WARP + lane;
+    const uint32_t c_out = bidn * THREADS_PER_WARP / 2 + lane;
+    constexpr uint32_t COL_STRIDE = Kernel_traits::CTAS * THREADS_PER_WARP;
+    for( uint32_t col = c, col_out = c_out; col < Kernel_traits::COLS; col += COL_STRIDE, col_out += COL_STRIDE / 2 ) {
+        // Each thread sums over NUM_ELT columns.
+        Vec<compute_t, NUM_ELT> dbeta_local, dgamma_local;
+        memset(&dgamma_local, 0, sizeof(dgamma_local));
+        memset(&dbeta_local, 0, sizeof(dbeta_local));
+        for( uint32_t row = warp; row < params.ctas_per_col; row += Kernel_traits::ROWS_PER_CTA ) {
+            index_t idx = row * Kernel_traits::COLS + col;
+
+            Vec<compute_t, NUM_ELT> dbeta_part, dgamma_part;
+            dbeta_part.load_from(params.dbeta_part, idx);
+            dgamma_part.load_from(params.dgamma_part, idx);
+            #pragma unroll
+            for( int it = 0; it < NUM_ELT; it++ ) {
+                dgamma_local.data.elt[it] += dgamma_part.data.elt[it];
+                dbeta_local.data.elt[it] += dbeta_part.data.elt[it];
+            }
+        }
+
+        void * smem_gamma = smem_;
+        void * smem_beta = &smem_[Kernel_traits::SMEM_BYTES_TRANSPOSE];
+
+        const int write_row = warp;
+        const int write_col = lane ^ write_row;
+        const int write_idx = write_row * THREADS_PER_WARP + write_col;
+
+        dgamma_local.store_to(smem_gamma, write_idx);
+        dbeta_local.store_to(smem_beta, write_idx);
+
+        __syncthreads();
+
+        // It would be probably safe to reuse the first row of smem_beta and smem_gamma
+        void * smem_gamma_out = &smem_[2 * Kernel_traits::SMEM_BYTES_TRANSPOSE];
+        void * smem_beta_out = &smem_[2 * Kernel_traits::SMEM_BYTES_TRANSPOSE + Kernel_traits::SMEM_BYTES_OUTPUT];
+
+
+        // More than one iter iff ROWS_PER_CTA < 32.
+        for( int w = warp; w < THREADS_PER_WARP; w += Kernel_traits::ROWS_PER_CTA ) {
+            const int read_row = lane;
+            const int read_col = w ^ read_row;
+            const int read_idx = read_row * THREADS_PER_WARP + read_col;
+
+            memset(&dbeta_local, 0, sizeof(dbeta_local));
+            memset(&dgamma_local, 0, sizeof(dgamma_local));
+
+            // Load beta and gamma transposed 
+            if(read_row < Kernel_traits::ROWS_PER_CTA){
+                dbeta_local.load_from(smem_beta, read_idx);
+                dgamma_local.load_from(smem_gamma, read_idx);
+            }
+
+            // Call reducer on the loaded value(s) and convert.
+            #pragma unroll
+            for( int it = 0; it < NUM_ELT; it++ ) {
+                compute_t b_i = dbeta_local.data.elt[it];
+                compute_t g_i = dgamma_local.data.elt[it];
+                b_i = reducer.allreduce(b_i, sum);
+                g_i = reducer.allreduce(g_i, sum);
+
+                dgamma_local.data.elt[it] = g_i;
+                dbeta_local.data.elt[it] = b_i;
+            }
+
+            // Leader stores the result at the current column.
+            if(lane == 0){
+                dgamma_local.store_to(smem_gamma_out, w);
+                dbeta_local.store_to(smem_beta_out, w);
+            }
+
+        }
+
+        // All writes done.
+        __syncthreads();
+
+        // Pack and store: 2-wide stores with half the threads.
+        if( warp == Kernel_traits::ROWS_PER_CTA - 1 && lane < THREADS_PER_WARP / 2 ) {
+
+            using src_t = typename TypeToVec2<compute_t>::Type;
+            using dst_t = typename TypeToVec2<weight_t>::Type;
+            Vec<src_t, NUM_ELT> dbeta_vec2, dgamma_vec2;
+            Vec<dst_t, NUM_ELT> dbeta_out2, dgamma_out2;
+
+            dgamma_vec2.load_from(smem_gamma_out, lane);
+            dbeta_vec2.load_from(smem_beta_out, lane);
+            #pragma unroll
+            for( int it = 0; it < NUM_ELT; it++ ) {
+                dgamma_out2.data.elt[it] = Converter<src_t,dst_t>::convert(dgamma_vec2.data.elt[it]);
+                dbeta_out2.data.elt[it] = Converter<src_t,dst_t>::convert(dbeta_vec2.data.elt[it]);
+            }
+            dgamma_out2.store_to(params.dgamma, col_out);
+            dbeta_out2.store_to(params.dbeta, col_out);
+
+        }
+    }
+}
+}  // namespace layer_norm
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_bwd_semi_cuda_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_bwd_semi_cuda_kernel.cu
new file mode 100644
index 0000000000..1908b045cd
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_bwd_semi_cuda_kernel.cu
@@ -0,0 +1,234 @@
+#include "ln.h"
+#include "ln_utils.cuh"
+#include "ln_kernel_traits.h"
+#include "ln_bwd_kernels.cuh"
+
+using namespace layer_norm;
+
+template<
+    typename weight_t,
+    typename input_t,
+    typename output_t,
+    typename compute_t,
+    typename index_t,
+    int HIDDEN_SIZE, 
+    int CTAS_PER_ROW, 
+    int WARPS_M, 
+    int WARPS_N, 
+    int BYTES_PER_LDG_MAIN,
+    int BYTES_PER_LDG_FINAL
+>
+void launch_(LaunchParams<BwdParams> &launch_params, const bool configure_params){
+
+    using Kernel_traits = Kernel_traits<weight_t,
+                                        input_t,
+                                        output_t,
+                                        compute_t,
+                                        index_t,
+                                        HIDDEN_SIZE,
+                                        CTAS_PER_ROW,
+                                        WARPS_M,
+                                        WARPS_N,
+                                        BYTES_PER_LDG_MAIN
+                                        >;
+    auto kernel = &ln_bwd_kernel<Kernel_traits>;
+
+    if( configure_params ) {
+        int ctas_per_sm;
+        cudaError status_ = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+            &ctas_per_sm, kernel, Kernel_traits::THREADS_PER_CTA, Kernel_traits::SMEM_BYTES);
+        launch_params.params.ctas_per_col = launch_params.props->multiProcessorCount * ctas_per_sm / Kernel_traits::CTAS_PER_ROW;
+        launch_params.barrier_size = 0;
+        launch_params.workspace_bytes = 0;
+        if(Kernel_traits::CTAS_PER_ROW > 1) {
+            launch_params.barrier_size = 2 * launch_params.params.ctas_per_col;
+            launch_params.workspace_bytes = launch_params.params.ctas_per_col 
+                                          * Kernel_traits::WARPS_M  
+                                          * Kernel_traits::CTAS_PER_ROW 
+                                          * sizeof(typename Kernel_traits::reduce_t)
+                                          * 2;
+        }
+        return;
+    }
+
+    if( Kernel_traits::SMEM_BYTES >= 48 * 1024 ) {
+        CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, Kernel_traits::SMEM_BYTES));
+    }
+    auto stream = launch_params.stream;
+    auto ctas_per_col = launch_params.params.ctas_per_col;
+
+    if( Kernel_traits::CTAS_PER_ROW == 1 ) {
+        kernel<<<ctas_per_col, Kernel_traits::THREADS_PER_CTA, Kernel_traits::SMEM_BYTES, stream>>>(launch_params.params);
+    } else {
+        dim3 grid(Kernel_traits::CTAS_PER_ROW * ctas_per_col);
+        dim3 block(Kernel_traits::THREADS_PER_CTA);
+        void *params_ = (void *)&launch_params.params;
+        cudaLaunchCooperativeKernel((void *)kernel, grid, block, (void **)&params_, Kernel_traits::SMEM_BYTES, stream);
+    }
+
+    using Kernel_traits_f = layer_norm::Kernel_traits_finalize<HIDDEN_SIZE,
+                                                               weight_t,
+                                                               input_t,
+                                                               output_t,
+                                                               compute_t,
+                                                               index_t,
+                                                               32 * 32,  // THREADS_PER_CTA
+                                                               BYTES_PER_LDG_FINAL>;
+
+    auto kernel_f = &layer_norm::ln_bwd_finalize_kernel<Kernel_traits_f>;
+    kernel_f<<<Kernel_traits_f::CTAS, Kernel_traits_f::THREADS_PER_CTA, 0, stream>>>(launch_params.params);
+}
+
+// Create backward launch function and register. Macro signature:
+//  HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, BYTES_PER_LDG_FINAL
+
+REGISTER_BWD_LAUNCHER(  768, fp32, fp32, fp32, fp32, 1, 4, 1, 16, 4);
+REGISTER_BWD_LAUNCHER(  768, fp16, fp16, fp16, fp32, 1, 4, 1, 16, 4);
+REGISTER_BWD_LAUNCHER(  768, fp16, fp32, fp16, fp32, 1, 4, 1, 16, 4);
+REGISTER_BWD_LAUNCHER(  768, bf16, bf16, bf16, fp32, 1, 4, 1, 16, 4);
+REGISTER_BWD_LAUNCHER(  768, bf16, fp32, bf16, fp32, 1, 4, 1, 16, 4);
+
+REGISTER_BWD_LAUNCHER( 1024, fp32, fp32, fp32, fp32, 1, 4, 1, 16, 4);
+REGISTER_BWD_LAUNCHER( 1024, fp16, fp16, fp16, fp32, 1, 4, 1, 16, 4);
+REGISTER_BWD_LAUNCHER( 1024, fp16, fp32, fp16, fp32, 1, 4, 1, 16, 4);
+REGISTER_BWD_LAUNCHER( 1024, bf16, bf16, bf16, fp32, 1, 4, 1, 16, 4);
+REGISTER_BWD_LAUNCHER( 1024, bf16, fp32, bf16, fp32, 1, 4, 1, 16, 4);
+
+REGISTER_BWD_LAUNCHER( 1536, fp32, fp32, fp32, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 1536, fp16, fp16, fp16, fp32, 1, 1, 4,  8, 4);
+REGISTER_BWD_LAUNCHER( 1536, fp16, fp32, fp16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 1536, bf16, bf16, bf16, fp32, 1, 1, 4,  8, 4);
+REGISTER_BWD_LAUNCHER( 1536, bf16, fp32, bf16, fp32, 1, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER( 2048, fp32, fp32, fp32, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 2048, fp16, fp16, fp16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 2048, fp16, fp32, fp16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 2048, bf16, bf16, bf16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 2048, bf16, fp32, bf16, fp32, 1, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER( 2304, fp32, fp32, fp32, fp32, 1, 1, 4,  8, 4);
+REGISTER_BWD_LAUNCHER( 2304, fp16, fp16, fp16, fp32, 1, 1, 4,  4, 4);
+REGISTER_BWD_LAUNCHER( 2304, fp16, fp32, fp16, fp32, 1, 1, 4,  8, 4);
+REGISTER_BWD_LAUNCHER( 2304, bf16, bf16, bf16, fp32, 1, 1, 4,  4, 4);
+REGISTER_BWD_LAUNCHER( 2304, bf16, fp32, bf16, fp32, 1, 1, 4,  8, 4);
+
+REGISTER_BWD_LAUNCHER( 3072, fp32, fp32, fp32, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 3072, fp16, fp16, fp16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 3072, fp16, fp32, fp16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 3072, bf16, bf16, bf16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 3072, bf16, fp32, bf16, fp32, 1, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER( 3840, fp32, fp32, fp32, fp32, 1, 1, 4, 8, 4);
+REGISTER_BWD_LAUNCHER( 3840, fp16, fp16, fp16, fp32, 1, 1, 4, 4, 4);
+REGISTER_BWD_LAUNCHER( 3840, fp16, fp32, fp16, fp32, 1, 1, 4, 8, 4);
+REGISTER_BWD_LAUNCHER( 3840, bf16, bf16, bf16, fp32, 1, 1, 4, 4, 4);
+REGISTER_BWD_LAUNCHER( 3840, bf16, fp32, bf16, fp32, 1, 1, 4, 8, 4);
+
+REGISTER_BWD_LAUNCHER( 4096, fp32, fp32, fp32, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 4096, fp16, fp16, fp16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 4096, fp16, fp32, fp16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 4096, bf16, bf16, bf16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 4096, bf16, fp32, bf16, fp32, 1, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER( 5120, fp32, fp32, fp32, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 5120, fp16, fp16, fp16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 5120, fp16, fp32, fp16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 5120, bf16, bf16, bf16, fp32, 1, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 5120, bf16, fp32, bf16, fp32, 1, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER( 6144, fp32, fp32, fp32, fp32, 1, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER( 6144, fp16, fp16, fp16, fp32, 1, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER( 6144, fp16, fp32, fp16, fp32, 1, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER( 6144, bf16, bf16, bf16, fp32, 1, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER( 6144, bf16, fp32, bf16, fp32, 1, 1, 8, 16, 4);
+
+REGISTER_BWD_LAUNCHER( 8192, fp32, fp32, fp32, fp32, 2, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 8192, fp16, fp16, fp16, fp32, 2, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 8192, fp16, fp32, fp16, fp32, 2, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 8192, bf16, bf16, bf16, fp32, 2, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER( 8192, bf16, fp32, bf16, fp32, 2, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER(10240, fp32, fp32, fp32, fp32, 2, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(10240, fp16, fp16, fp16, fp32, 2, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(10240, fp16, fp32, fp16, fp32, 2, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(10240, bf16, bf16, bf16, fp32, 2, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(10240, bf16, fp32, bf16, fp32, 2, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER(12288, fp32, fp32, fp32, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(12288, fp16, fp16, fp16, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(12288, fp16, fp32, fp16, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(12288, bf16, bf16, bf16, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(12288, bf16, fp32, bf16, fp32, 4, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER(12800, fp32, fp32, fp32, fp32, 5, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(12800, fp16, fp16, fp16, fp32, 5, 1, 4,  8, 4);
+REGISTER_BWD_LAUNCHER(12800, fp16, fp32, fp16, fp32, 5, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(12800, bf16, bf16, bf16, fp32, 5, 1, 4,  8, 4);
+REGISTER_BWD_LAUNCHER(12800, bf16, fp32, bf16, fp32, 5, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER(15360, fp32, fp32, fp32, fp32, 4, 1, 4,  8, 4);
+REGISTER_BWD_LAUNCHER(15360, fp16, fp16, fp16, fp32, 4, 1, 4,  4, 4);
+REGISTER_BWD_LAUNCHER(15360, fp16, fp32, fp16, fp32, 4, 1, 4,  8, 4);
+REGISTER_BWD_LAUNCHER(15360, bf16, bf16, bf16, fp32, 4, 1, 4,  4, 4);
+REGISTER_BWD_LAUNCHER(15360, bf16, fp32, bf16, fp32, 4, 1, 4,  8, 4);
+
+REGISTER_BWD_LAUNCHER(16384, fp32, fp32, fp32, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(16384, fp16, fp16, fp16, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(16384, fp16, fp32, fp16, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(16384, bf16, bf16, bf16, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(16384, bf16, fp32, bf16, fp32, 4, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER(18432, fp32, fp32, fp32, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(18432, fp16, fp16, fp16, fp32, 4, 1, 4,  8, 4);
+REGISTER_BWD_LAUNCHER(18432, fp16, fp32, fp16, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(18432, bf16, bf16, bf16, fp32, 4, 1, 4,  8, 4);
+REGISTER_BWD_LAUNCHER(18432, bf16, fp32, bf16, fp32, 4, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER(20480, fp32, fp32, fp32, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(20480, fp16, fp16, fp16, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(20480, fp16, fp32, fp16, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(20480, bf16, bf16, bf16, fp32, 4, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(20480, bf16, fp32, bf16, fp32, 4, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER(24576, fp32, fp32, fp32, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(24576, fp16, fp16, fp16, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(24576, fp16, fp32, fp16, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(24576, bf16, bf16, bf16, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(24576, bf16, fp32, bf16, fp32, 4, 1, 8, 16, 4);
+
+REGISTER_BWD_LAUNCHER(25600, fp32, fp32, fp32, fp32, 5, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(25600, fp16, fp16, fp16, fp32, 5, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(25600, fp16, fp32, fp16, fp32, 5, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(25600, bf16, bf16, bf16, fp32, 5, 1, 4, 16, 4);
+REGISTER_BWD_LAUNCHER(25600, bf16, fp32, bf16, fp32, 5, 1, 4, 16, 4);
+
+REGISTER_BWD_LAUNCHER(30720, fp32, fp32, fp32, fp32, 4, 1, 8, 8, 4);
+REGISTER_BWD_LAUNCHER(30720, fp16, fp16, fp16, fp32, 4, 1, 8, 4, 4);
+REGISTER_BWD_LAUNCHER(30720, fp16, fp32, fp16, fp32, 4, 1, 8, 8, 4);
+REGISTER_BWD_LAUNCHER(30720, bf16, bf16, bf16, fp32, 4, 1, 8, 4, 4);
+REGISTER_BWD_LAUNCHER(30720, bf16, fp32, bf16, fp32, 4, 1, 8, 8, 4);
+
+REGISTER_BWD_LAUNCHER(32768, fp32, fp32, fp32, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(32768, fp16, fp16, fp16, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(32768, fp16, fp32, fp16, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(32768, bf16, bf16, bf16, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(32768, bf16, fp32, bf16, fp32, 4, 1, 8, 16, 4);
+
+REGISTER_BWD_LAUNCHER(40960, fp32, fp32, fp32, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(40960, fp16, fp16, fp16, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(40960, fp16, fp32, fp16, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(40960, bf16, bf16, bf16, fp32, 4, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(40960, bf16, fp32, bf16, fp32, 4, 1, 8, 16, 4);
+
+REGISTER_BWD_LAUNCHER(49152, fp32, fp32, fp32, fp32, 8, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(49152, fp16, fp16, fp16, fp32, 8, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(49152, fp16, fp32, fp16, fp32, 8, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(49152, bf16, bf16, bf16, fp32, 8, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(49152, bf16, fp32, bf16, fp32, 8, 1, 8, 16, 4);
+
+REGISTER_BWD_LAUNCHER(65536, fp32, fp32, fp32, fp32, 8, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(65536, fp16, fp16, fp16, fp32, 8, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(65536, fp16, fp32, fp16, fp32, 8, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(65536, bf16, bf16, bf16, fp32, 8, 1, 8, 16, 4);
+REGISTER_BWD_LAUNCHER(65536, bf16, fp32, bf16, fp32, 8, 1, 8, 16, 4);
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_fwd_cuda_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_fwd_cuda_kernel.cu
new file mode 100644
index 0000000000..bc29c3518a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_fwd_cuda_kernel.cu
@@ -0,0 +1,222 @@
+#include "ln.h"
+#include "ln_utils.cuh"
+#include "ln_kernel_traits.h"
+#include "ln_fwd_kernels.cuh"
+
+using namespace layer_norm;
+
+template<
+    typename weight_t,
+    typename input_t,
+    typename output_t,
+    typename compute_t,
+    typename index_t,
+    int HIDDEN_SIZE, 
+    int CTAS_PER_ROW, 
+    int WARPS_M, 
+    int WARPS_N, 
+    int BYTES_PER_LDG
+>
+void launch_(LaunchParams<FwdParams> &launch_params, const bool configure_params){
+
+    using Kernel_traits = Kernel_traits<weight_t,
+                                        input_t,
+                                        output_t,
+                                        compute_t,
+                                        index_t,
+                                        HIDDEN_SIZE,
+                                        CTAS_PER_ROW,
+                                        WARPS_M,
+                                        WARPS_N,
+                                        BYTES_PER_LDG
+                                        >;
+    auto kernel = &ln_fwd_kernel<Kernel_traits>;
+
+    if( configure_params ) {
+        int ctas_per_sm;
+        cudaError status_ = cudaOccupancyMaxActiveBlocksPerMultiprocessor(
+            &ctas_per_sm, kernel, Kernel_traits::THREADS_PER_CTA, Kernel_traits::SMEM_BYTES_FWD);
+        launch_params.params.ctas_per_col = launch_params.props->multiProcessorCount * ctas_per_sm / Kernel_traits::CTAS_PER_ROW;
+        launch_params.barrier_size = 0;
+        launch_params.workspace_bytes = 0;
+        if(Kernel_traits::CTAS_PER_ROW > 1) {
+            launch_params.barrier_size = 2 * launch_params.params.ctas_per_col;
+            launch_params.workspace_bytes = launch_params.params.ctas_per_col 
+                                          * Kernel_traits::WARPS_M  
+                                          * Kernel_traits::CTAS_PER_ROW 
+                                          * sizeof(typename Kernel_traits::Stats::stats_t)
+                                          * 2;
+        }
+        return;
+    }
+
+    if( Kernel_traits::SMEM_BYTES_FWD >= 48 * 1024 ) {
+        CHECK_CUDA(cudaFuncSetAttribute(kernel, cudaFuncAttributeMaxDynamicSharedMemorySize, Kernel_traits::SMEM_BYTES_FWD));
+    }
+    auto stream = launch_params.stream;
+    auto ctas_per_col = launch_params.params.ctas_per_col;
+
+    if( Kernel_traits::CTAS_PER_ROW == 1 ) {
+        kernel<<<ctas_per_col, Kernel_traits::THREADS_PER_CTA, Kernel_traits::SMEM_BYTES_FWD, stream>>>(launch_params.params);
+    } else {
+        dim3 grid(Kernel_traits::CTAS_PER_ROW * ctas_per_col);
+        dim3 block(Kernel_traits::THREADS_PER_CTA);
+        void *params_ = (void *)&launch_params.params;
+        cudaLaunchCooperativeKernel((void *)kernel, grid, block, (void **)&params_, Kernel_traits::SMEM_BYTES_FWD, stream);
+    }
+
+}
+
+// Create forward launch function and register. Macro signature:
+//  HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG
+
+REGISTER_FWD_LAUNCHER(  768, fp32, fp32, fp32, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER(  768, fp16, fp16, fp16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER(  768, fp16, fp32, fp16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER(  768, bf16, bf16, bf16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER(  768, bf16, fp32, bf16, fp32, 1, 4, 1, 16);
+
+REGISTER_FWD_LAUNCHER( 1024, fp32, fp32, fp32, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 1024, fp16, fp16, fp16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 1024, fp16, fp32, fp16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 1024, bf16, bf16, bf16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 1024, bf16, fp32, bf16, fp32, 1, 4, 1, 16);
+
+REGISTER_FWD_LAUNCHER( 1536, fp32, fp32, fp32, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 1536, fp16, fp16, fp16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 1536, fp16, fp32, fp16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 1536, bf16, bf16, bf16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 1536, bf16, fp32, bf16, fp32, 1, 4, 1, 16);
+
+REGISTER_FWD_LAUNCHER( 2048, fp32, fp32, fp32, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 2048, fp16, fp16, fp16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 2048, fp16, fp32, fp16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 2048, bf16, bf16, bf16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 2048, bf16, fp32, bf16, fp32, 1, 4, 1, 16);
+
+REGISTER_FWD_LAUNCHER( 2304, fp32, fp32, fp32, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 2304, fp16, fp16, fp16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 2304, fp16, fp32, fp16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 2304, bf16, bf16, bf16, fp32, 1, 4, 1, 16);
+REGISTER_FWD_LAUNCHER( 2304, bf16, fp32, bf16, fp32, 1, 4, 1, 16);
+
+REGISTER_FWD_LAUNCHER( 3072, fp32, fp32, fp32, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 3072, fp16, fp16, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 3072, fp16, fp32, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 3072, bf16, bf16, bf16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 3072, bf16, fp32, bf16, fp32, 1, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER( 3840, fp32, fp32, fp32, fp32, 1, 1, 4,  4);
+REGISTER_FWD_LAUNCHER( 3840, fp16, fp16, fp16, fp32, 1, 1, 4,  4);
+REGISTER_FWD_LAUNCHER( 3840, fp16, fp32, fp16, fp32, 1, 1, 4,  4);
+REGISTER_FWD_LAUNCHER( 3840, bf16, bf16, bf16, fp32, 1, 1, 4,  4);
+REGISTER_FWD_LAUNCHER( 3840, bf16, fp32, bf16, fp32, 1, 1, 4,  4);
+
+REGISTER_FWD_LAUNCHER( 4096, fp32, fp32, fp32, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 4096, fp16, fp16, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 4096, fp16, fp32, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 4096, bf16, bf16, bf16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 4096, bf16, fp32, bf16, fp32, 1, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER( 5120, fp32, fp32, fp32, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 5120, fp16, fp16, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 5120, fp16, fp32, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 5120, bf16, bf16, bf16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 5120, bf16, fp32, bf16, fp32, 1, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER( 6144, fp32, fp32, fp32, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 6144, fp16, fp16, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 6144, fp16, fp32, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 6144, bf16, bf16, bf16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 6144, bf16, fp32, bf16, fp32, 1, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER( 8192, fp32, fp32, fp32, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 8192, fp16, fp16, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 8192, fp16, fp32, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 8192, bf16, bf16, bf16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER( 8192, bf16, fp32, bf16, fp32, 1, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER(10240, fp32, fp32, fp32, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(10240, fp16, fp16, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(10240, fp16, fp32, fp16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(10240, bf16, bf16, bf16, fp32, 1, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(10240, bf16, fp32, bf16, fp32, 1, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER(12288, fp32, fp32, fp32, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(12288, fp16, fp16, fp16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(12288, fp16, fp32, fp16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(12288, bf16, bf16, bf16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(12288, bf16, fp32, bf16, fp32, 2, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER(12800, fp32, fp32, fp32, fp32, 2, 1, 4,  4);
+REGISTER_FWD_LAUNCHER(12800, fp16, fp16, fp16, fp32, 2, 1, 4,  4);
+REGISTER_FWD_LAUNCHER(12800, fp16, fp32, fp16, fp32, 2, 1, 4,  4);
+REGISTER_FWD_LAUNCHER(12800, bf16, bf16, bf16, fp32, 2, 1, 4,  4);
+REGISTER_FWD_LAUNCHER(12800, bf16, fp32, bf16, fp32, 2, 1, 4,  4);
+
+REGISTER_FWD_LAUNCHER(15360, fp32, fp32, fp32, fp32, 2, 1, 4,  8);
+REGISTER_FWD_LAUNCHER(15360, fp16, fp16, fp16, fp32, 2, 1, 4,  8);
+REGISTER_FWD_LAUNCHER(15360, fp16, fp32, fp16, fp32, 2, 1, 4,  8);
+REGISTER_FWD_LAUNCHER(15360, bf16, bf16, bf16, fp32, 2, 1, 4,  8);
+REGISTER_FWD_LAUNCHER(15360, bf16, fp32, bf16, fp32, 2, 1, 4,  8);
+
+REGISTER_FWD_LAUNCHER(16384, fp32, fp32, fp32, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(16384, fp16, fp16, fp16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(16384, fp16, fp32, fp16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(16384, bf16, bf16, bf16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(16384, bf16, fp32, bf16, fp32, 2, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER(18432, fp32, fp32, fp32, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(18432, fp16, fp16, fp16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(18432, fp16, fp32, fp16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(18432, bf16, bf16, bf16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(18432, bf16, fp32, bf16, fp32, 4, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER(20480, fp32, fp32, fp32, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(20480, fp16, fp16, fp16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(20480, fp16, fp32, fp16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(20480, bf16, bf16, bf16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(20480, bf16, fp32, bf16, fp32, 2, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER(24576, fp32, fp32, fp32, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(24576, fp16, fp16, fp16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(24576, fp16, fp32, fp16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(24576, bf16, bf16, bf16, fp32, 2, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(24576, bf16, fp32, bf16, fp32, 2, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER(25600, fp32, fp32, fp32, fp32, 4, 1, 4,  4);
+REGISTER_FWD_LAUNCHER(25600, fp16, fp16, fp16, fp32, 2, 1, 4,  8);
+REGISTER_FWD_LAUNCHER(25600, fp16, fp32, fp16, fp32, 4, 1, 4,  4);
+REGISTER_FWD_LAUNCHER(25600, bf16, bf16, bf16, fp32, 2, 1, 4,  8);
+REGISTER_FWD_LAUNCHER(25600, bf16, fp32, bf16, fp32, 4, 1, 4,  4);
+
+REGISTER_FWD_LAUNCHER(30720, fp32, fp32, fp32, fp32, 4, 1, 4,  4);
+REGISTER_FWD_LAUNCHER(30720, fp16, fp16, fp16, fp32, 4, 1, 4,  4);
+REGISTER_FWD_LAUNCHER(30720, fp16, fp32, fp16, fp32, 4, 1, 4,  4);
+REGISTER_FWD_LAUNCHER(30720, bf16, bf16, bf16, fp32, 4, 1, 4,  4);
+REGISTER_FWD_LAUNCHER(30720, bf16, fp32, bf16, fp32, 4, 1, 4,  4);
+
+REGISTER_FWD_LAUNCHER(32768, fp32, fp32, fp32, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(32768, fp16, fp16, fp16, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(32768, fp16, fp32, fp16, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(32768, bf16, bf16, bf16, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(32768, bf16, fp32, bf16, fp32, 4, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER(40960, fp32, fp32, fp32, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(40960, fp16, fp16, fp16, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(40960, fp16, fp32, fp16, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(40960, bf16, bf16, bf16, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(40960, bf16, fp32, bf16, fp32, 4, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER(49152, fp32, fp32, fp32, fp32, 8, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(49152, fp16, fp16, fp16, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(49152, fp16, fp32, fp16, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(49152, bf16, bf16, bf16, fp32, 4, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(49152, bf16, fp32, bf16, fp32, 4, 1, 4, 16);
+
+REGISTER_FWD_LAUNCHER(65536, fp32, fp32, fp32, fp32, 8, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(65536, fp16, fp16, fp16, fp32, 8, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(65536, fp16, fp32, fp16, fp32, 8, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(65536, bf16, bf16, bf16, fp32, 8, 1, 4, 16);
+REGISTER_FWD_LAUNCHER(65536, bf16, fp32, bf16, fp32, 8, 1, 4, 16);
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_fwd_kernels.cuh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_fwd_kernels.cuh
new file mode 100644
index 0000000000..64e72974f5
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_fwd_kernels.cuh
@@ -0,0 +1,110 @@
+#pragma once
+
+#include "ln.h"
+
+namespace layer_norm {
+
+template<typename Ktraits>
+__global__ __launch_bounds__(Ktraits::THREADS_PER_CTA) 
+void ln_fwd_kernel(FwdParams params) {
+
+    enum { ROWS_PER_CTA = Ktraits::ROWS_PER_CTA };
+    enum { WARPS_N = Ktraits::WARPS_N };
+    enum { WARPS_M = Ktraits::WARPS_M };
+    enum { THREADS_PER_ROW = Ktraits::THREADS_PER_ROW };
+    enum { VEC_COLS_PER_LDG = Ktraits::VEC_COLS_PER_LDG };
+    enum { BYTES_PER_ROW = Ktraits::BYTES_PER_ROW };
+    enum { LDGS = Ktraits::LDGS };
+    enum { NUM_ELTS = Ktraits::NUM_ELTS };
+    enum { CTAS_PER_ROW = Ktraits::CTAS_PER_ROW };
+
+    using output_t = typename Ktraits::output_t;
+    using index_t = typename Ktraits::index_t;
+    using compute_t = typename Ktraits::compute_t;
+    using Ivec = typename Ktraits::Ivec;
+    using Ovec = typename Ktraits::Ovec;
+    using Wvec = typename Ktraits::Wvec;
+    using Cvec = typename Ktraits::Cvec;
+
+    using Stats = typename Ktraits::Stats;
+    using stats_t = typename Stats::stats_t;
+
+    extern __shared__ char smem_[];
+
+    const index_t tidx = threadIdx.x;
+    const index_t bidn = blockIdx.x % CTAS_PER_ROW;
+    const index_t bidm = blockIdx.x / CTAS_PER_ROW;
+    const index_t lane = tidx % THREADS_PER_WARP;
+    const index_t warp = tidx / THREADS_PER_WARP;
+    const index_t warp_m = warp / WARPS_N;
+    const index_t warp_n = warp % WARPS_N;
+
+    const index_t r = bidm * ROWS_PER_CTA + warp_m;
+    const index_t c = bidn * THREADS_PER_ROW + warp_n * THREADS_PER_WARP + lane;
+
+    Stats stats(params, bidm, bidn, warp_m, warp_n, lane, smem_);
+
+    compute_t *mu_ptr = static_cast<compute_t *>(params.mu);
+    compute_t *rs_ptr = static_cast<compute_t *>(params.rs);
+
+    Wvec gamma[LDGS];
+    Wvec beta[LDGS];
+    index_t idx = c;
+    #pragma unroll
+    for( int it = 0; it < LDGS; it++ ) {
+        gamma[it].load_from(params.gamma, idx);
+        beta[it].load_from(params.beta, idx);
+        idx += VEC_COLS_PER_LDG;
+    }
+
+    constexpr compute_t rn = 1.f / compute_t(Ktraits::COLS);
+
+    for( int row = r; row < params.rows; row += params.ctas_per_col * ROWS_PER_CTA ) {
+        Ivec x[LDGS];
+        index_t idx = row * Ktraits::VEC_COLS + c;
+        compute_t xf[LDGS * NUM_ELTS];
+        #pragma unroll
+        for( int it = 0; it < LDGS; it++ ) {
+            x[it].load_from(params.x, idx);
+            #pragma unroll
+            for( int jt = 0; jt < NUM_ELTS; jt++ ) {
+                compute_t x_ij = compute_t(x[it].data.elt[jt]);
+                xf[it * NUM_ELTS + jt] =  x_ij;
+            }
+            idx += VEC_COLS_PER_LDG;
+        }
+
+        stats_t s = stats.compute(xf, rn);
+
+        compute_t mu = layer_norm::Get<0>::of<stats_t, compute_t>(s);
+        compute_t m2 = layer_norm::Get<1>::of<stats_t, compute_t>(s);
+
+        if( bidn == 0 && warp_n == 0 && lane == 0 ) {
+            mu_ptr[row] = mu;
+        }
+
+        compute_t rs = rsqrtf(rn * m2 + params.epsilon);
+
+        if( bidn == 0 && warp_n == 0 && lane == 0 ) {
+            rs_ptr[row] = rs;
+        }
+
+        Ovec z[LDGS];
+        idx = row * Ktraits::VEC_COLS + c;
+        #pragma unroll
+        for( int it = 0; it < LDGS; it++ ) {
+            #pragma unroll
+            for( int jt = 0; jt < NUM_ELTS; jt++ ) {
+                output_t y_ij = output_t(rs * (xf[it * NUM_ELTS + jt] - mu));
+                output_t g_ij = gamma[it].data.elt[jt];
+                output_t b_ij = beta[it].data.elt[jt];
+                z[it].data.elt[jt] = (g_ij * y_ij + b_ij);
+            }
+            z[it].store_to(params.z, idx);
+            idx += VEC_COLS_PER_LDG;
+        }
+
+    }
+}
+
+}  // namespace layer_norm
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_kernel_traits.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_kernel_traits.h
new file mode 100644
index 0000000000..ed745c5eec
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_kernel_traits.h
@@ -0,0 +1,159 @@
+#pragma once
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace layer_norm {
+template<
+    uint32_t HIDDEN_SIZE_,
+    typename weight_t_,
+    typename input_t_,
+    typename output_t_,
+    typename compute_t_,
+    typename index_t_,
+    uint32_t THREADS_PER_CTA_
+>
+struct Kernel_traits_base {
+
+    using weight_t = weight_t_;
+    using input_t = input_t_;
+    using output_t = output_t_;
+    using compute_t = compute_t_;
+    using index_t = index_t_;
+
+    enum { HIDDEN_SIZE = HIDDEN_SIZE_ };
+    enum { THREADS_PER_CTA = THREADS_PER_CTA_ };
+    enum { THREADS_PER_WARP = 32 };
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<
+    uint32_t HIDDEN_SIZE_,
+    typename weight_t_,
+    typename input_t_,
+    typename output_t_,
+    typename compute_t_,
+    typename index_t_,
+    uint32_t THREADS_PER_CTA_,
+    uint32_t BYTES_PER_LDG_,
+    typename Base = Kernel_traits_base<HIDDEN_SIZE_,
+                                        weight_t_,
+                                        input_t_,
+                                        output_t_,
+                                        compute_t_,
+                                        index_t_,
+                                        THREADS_PER_CTA_>
+>
+struct Kernel_traits_finalize : public Base {
+    enum { ROWS_PER_CTA = Base::THREADS_PER_CTA / Base::THREADS_PER_WARP };
+    static_assert((int) ROWS_PER_CTA <= (int) Base::THREADS_PER_WARP);
+    // Bytes per global load from the input. 
+    enum { BYTES_PER_LDG = BYTES_PER_LDG_ };
+    // Number of elements fetched by a global load.
+    enum { ELTS_PER_LDG = BYTES_PER_LDG / sizeof(compute_t_) };
+    // Bytes per global store of the weights.
+    enum { BYTES_PER_STG = ELTS_PER_LDG * sizeof(weight_t_) };
+    static_assert(sizeof(BYTES_PER_LDG) == 4, "Conflict-free smem transpose only implemented for 4B compute type!");
+    static_assert(Base::THREADS_PER_CTA == ROWS_PER_CTA * Base::THREADS_PER_WARP, "We assume one warp per row!");
+    // The total number of BYTES_PER_LDG-wide words in a hidden vector.
+    enum { COLS = HIDDEN_SIZE_ * sizeof(compute_t_) / BYTES_PER_LDG };
+    static_assert(COLS * BYTES_PER_LDG == HIDDEN_SIZE_ * sizeof(compute_t_));
+
+    // Shared memory size to transpose the CTA result.
+    enum { SMEM_BYTES_TRANSPOSE = Base::THREADS_PER_CTA * BYTES_PER_LDG };
+    // Shared memory size to coalsece the CTA result.
+    enum { SMEM_BYTES_OUTPUT = Base::THREADS_PER_WARP * BYTES_PER_LDG };
+    // Shared memory requirement per CTA. 
+    enum { SMEM_BYTES_PER_CTA = 2 * SMEM_BYTES_TRANSPOSE + 2 * SMEM_BYTES_OUTPUT };
+
+    // The type of the reducer.
+    using Reducer = layer_norm::Reducer<compute_t_, 1, 1, 1>;
+
+    // Condition for the whole CTA to participate in syncthreads.
+    static_assert(COLS % Base::THREADS_PER_WARP == 0);
+    enum { CTAS = COLS / Base::THREADS_PER_WARP };
+}; 
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+
+template<
+    typename weight_t_,
+    typename input_t_,
+    typename output_t_,
+    typename compute_t_,
+    typename index_t_,
+    uint32_t HIDDEN_SIZE_, 
+    uint32_t CTAS_PER_ROW_, 
+    uint32_t WARPS_M_, 
+    uint32_t WARPS_N_, 
+    uint32_t BYTES_PER_LDG_ = 16,
+    typename Base = Kernel_traits_base<
+        HIDDEN_SIZE_,
+        weight_t_, 
+        input_t_, 
+        output_t_, 
+        compute_t_, 
+        index_t_, 
+        WARPS_M_*WARPS_N_*THREADS_PER_WARP
+        >
+>
+struct Kernel_traits : public Base {
+
+    using input_t = typename Base::input_t;
+    using weight_t = typename Base::weight_t;
+    using compute_t = typename Base::compute_t;
+    using output_t = typename Base::output_t;
+    using index_t = typename Base::index_t;
+
+    enum { CTAS_PER_ROW = CTAS_PER_ROW_ };
+    enum { WARPS_M = WARPS_M_ };
+    enum { WARPS_N = WARPS_N_ };
+    enum { COLS = HIDDEN_SIZE_ };
+    enum { HIDDEN_SIZE = HIDDEN_SIZE_ };
+    enum { BYTES_PER_LDG = BYTES_PER_LDG_ };
+    enum { NUM_ELTS = BYTES_PER_LDG / sizeof(input_t) };
+
+    enum { THREADS_PER_ROW = WARPS_N * THREADS_PER_WARP };
+    enum { THREADS_PER_CTA = WARPS_M * THREADS_PER_ROW };
+    enum { ROWS_PER_CTA = WARPS_M };
+
+    enum { BYTES_PER_ROW = COLS * sizeof(input_t) };
+    enum { BYTES_PER_ROW_PER_CTA = THREADS_PER_ROW * BYTES_PER_LDG };
+    // Multi-row per CTA not supported for multi-CTA => no smem for WGRAD needed
+    enum { SMEM_BYTES_WGRAD = CTAS_PER_ROW > 1 ? 0 : ROWS_PER_CTA * COLS * sizeof(compute_t) };
+    static_assert(WARPS_M == 1 || CTAS_PER_ROW == 1);
+
+    using reduce_t = typename layer_norm::TypeToVec2<compute_t>::Type;
+    using Reducer = layer_norm::Reducer<reduce_t, CTAS_PER_ROW, WARPS_M, WARPS_N>; 
+
+    enum { SMEM_BYTES_DGRAD = Reducer::SMEM_BYTES };
+    enum { SMEM_BYTES = SMEM_BYTES_DGRAD  + SMEM_BYTES_WGRAD };
+
+    using Ivec = layer_norm::Vec<input_t, NUM_ELTS>;
+    using Ovec = layer_norm::Vec<output_t, NUM_ELTS>;
+    using Wvec = layer_norm::Vec<weight_t, NUM_ELTS>;
+    using Cvec = layer_norm::Vec<compute_t, NUM_ELTS>;
+    enum { ELTS_PER_LDG = BYTES_PER_LDG / sizeof(input_t) };
+
+    // Assume that each thread can handle the same number of elements in the output and weights as in the input.
+    static_assert(sizeof(input_t) >= sizeof(output_t));
+    static_assert(sizeof(input_t) >= sizeof(weight_t));
+    // The number of columns fetched per load from input: one per thread.
+    enum { VEC_COLS_PER_LDG =  CTAS_PER_ROW * THREADS_PER_ROW };
+    // The total number of vectorized loads/stores per hidden vector.
+    enum { VEC_COLS = COLS / ELTS_PER_LDG };
+    // The number of loads per thread for the input.
+    enum { LDGS = VEC_COLS / VEC_COLS_PER_LDG };
+    static_assert(LDGS * VEC_COLS_PER_LDG  == VEC_COLS);
+    //static_assert(LDGS * BYTES_PER_ROW_PER_CTA * CTAS_PER_ROW == BYTES_PER_ROW, "");
+
+    using Stats = layer_norm::Stats<compute_t, CTAS_PER_ROW, WARPS_M, WARPS_N>;
+    enum { SMEM_BYTES_FWD = Stats::SMEM_BYTES };
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace layer_norm
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_utils.cuh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_utils.cuh
new file mode 100644
index 0000000000..e18d36de78
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/layer_norm/ln_utils.cuh
@@ -0,0 +1,733 @@
+#pragma once
+
+#include <cassert>
+
+#include <cuda_bf16.h>
+#include <cuda_fp16.h>
+
+#include "ln.h"
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+constexpr uint32_t THREADS_PER_WARP = 32;
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline void check_cuda_(cudaError_t status, const char *file, int line) {
+    if( status != cudaSuccess ) {
+        fprintf(stderr, "CUDA Error: %s %s %d\n", cudaGetErrorString(status), file, line);
+        exit(status);
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define CHECK_CUDA(ans)                                                                                                        \
+    { check_cuda_((ans), __FILE__, __LINE__); }
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define DIVUP(x, y) (((x) + ((y)-1)) / (y))
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define REGISTER_FWD_LAUNCHER(HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG)          \
+    void ln_fwd_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE(LaunchParams<FwdParams> &launch_params,                  \
+                                                                      const bool configure_params) {                           \
+        launch_<WTYPE, ITYPE, OTYPE, CTYPE, uint32_t, HIDDEN_SIZE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG>(             \
+            launch_params, configure_params);                                                                                  \
+    }                                                                                                                          \
+    static FwdRegistrar<WTYPE, ITYPE, OTYPE, CTYPE, HIDDEN_SIZE> reg_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE(    \
+        ln_fwd_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE)
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+#define REGISTER_BWD_LAUNCHER(                                                                                                 \
+    HIDDEN_SIZE, WTYPE, ITYPE, OTYPE, CTYPE, CTAS_PER_ROW, WARPS_M, WARPS_N, BYTES_PER_LDG, BYTES_PER_LDG_FINALIZE)            \
+    void ln_bwd_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE(LaunchParams<BwdParams> &launch_params,                  \
+                                                                      const bool configure_params) {                           \
+        launch_<WTYPE,                                                                                                         \
+                ITYPE,                                                                                                         \
+                OTYPE,                                                                                                         \
+                CTYPE,                                                                                                         \
+                uint32_t,                                                                                                      \
+                HIDDEN_SIZE,                                                                                                   \
+                CTAS_PER_ROW,                                                                                                  \
+                WARPS_M,                                                                                                       \
+                WARPS_N,                                                                                                       \
+                BYTES_PER_LDG,                                                                                                 \
+                BYTES_PER_LDG_FINALIZE>(launch_params, configure_params);                                                      \
+    }                                                                                                                          \
+    static BwdRegistrar<WTYPE, ITYPE, OTYPE, CTYPE, HIDDEN_SIZE> reg_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE(    \
+        ln_bwd_##HIDDEN_SIZE##_##WTYPE##_##ITYPE##_##OTYPE##_##CTYPE)
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ float2 operator+(const float2 & a, const float2 & b){
+    return {a.x + b.x, a.y + b.y};
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+inline __device__ void operator+=(float2 & a, const float2 & b){
+    a.x += b.x;
+    a.y += b.y;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+struct Sum {
+    inline __device__ Sum(){}
+    inline __device__ T operator()(const T &a, const T &b){
+        return a + b;
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+inline __device__ T warp_shuffle_xor(const T & x, uint32_t idx){
+    return __shfl_xor_sync(uint32_t(-1), x, idx);
+}
+
+template<>
+inline __device__ float2 warp_shuffle_xor<float2>(const float2 & x, uint32_t idx){
+    return { warp_shuffle_xor(x.x, idx), warp_shuffle_xor(x.y, idx) };
+}
+
+template<typename T>
+inline __device__ T warp_shuffle_down(const T & x, uint32_t idx){
+    return __shfl_down_sync(uint32_t(-1), x, idx);
+}
+
+template<>
+inline __device__ float2 warp_shuffle_down<float2>(const float2 & x, uint32_t idx){
+    return { warp_shuffle_down(x.x, idx), warp_shuffle_down(x.y, idx) };
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+namespace layer_norm {
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct uint16 {
+    uint4 u;
+    uint4 v;
+    uint4 s;
+    uint4 t;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+struct uint8 {
+    uint4 u;
+    uint4 v;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int BYTES>
+struct BytesToType {};
+
+template<>
+struct BytesToType<64> {
+    using Type = uint16;
+    static_assert(sizeof(Type) == 64);
+};
+
+template<>
+struct BytesToType<32> {
+    using Type = uint8;
+    static_assert(sizeof(Type) == 32);
+};
+
+template<>
+struct BytesToType<16> {
+    using Type = uint4;
+    static_assert(sizeof(Type) == 16);
+};
+
+template<>
+struct BytesToType<8> {
+    using Type = uint64_t;
+    static_assert(sizeof(Type) == 8);
+};
+
+template<>
+struct BytesToType<4> {
+    using Type = uint32_t;
+    static_assert(sizeof(Type) == 4);
+};
+
+template<>
+struct BytesToType<2> {
+    using Type = uint16_t;
+    static_assert(sizeof(Type) == 2);
+};
+
+template<>
+struct BytesToType<1> {
+    using Type = uint8_t;
+    static_assert(sizeof(Type) == 1);
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+struct TypeToVec2 {};
+
+template<>
+struct TypeToVec2<float> {
+    using Type = float2;
+};
+
+template<>
+struct TypeToVec2<half> {
+    using Type = half2;
+};
+
+template<>
+struct TypeToVec2<nv_bfloat16> {
+    using Type = nv_bfloat162;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<int INDEX>
+struct Get {
+    template<typename T, typename R>
+    static inline __device__ R of(const T &vec);
+};
+
+template<>
+template<typename T, typename R>
+inline __device__ R Get<0>::of(const T &vec) {
+    return vec.x;
+}
+
+template<>
+template<typename T, typename R>
+inline __device__ R Get<1>::of(const T &vec) {
+    return vec.y;
+}
+
+template<>
+template<typename T, typename R>
+inline __device__ R Get<2>::of(const T &vec) {
+    return vec.z;
+}
+
+template<>
+template<typename T, typename R>
+inline __device__ R Get<3>::of(const T &vec) {
+    return vec.w;
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Src, typename Dst>
+struct Converter{
+    static inline __device__ Dst convert(const Src &from) {
+        return Dst(from);
+    }
+};
+
+template<>
+struct Converter<float2, half2>{
+    static inline __device__ half2 convert(const float2 &x) {
+        return __float22half2_rn(x);
+    }
+};
+
+template<>
+struct Converter<float2, nv_bfloat162>{
+    static inline __device__ nv_bfloat162 convert(const float2 &x) {
+#if __CUDA_ARCH__ >= 800
+        return __float22bfloat162_rn(x);
+#else
+        union {
+            nv_bfloat162 raw;
+            nv_bfloat16 x;
+            nv_bfloat16 y;
+        } tmp;
+        tmp.x = __float2bfloat16_rn(x.x);
+        tmp.y = __float2bfloat16_rn(x.y);
+        return tmp.raw;
+#endif
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T>
+struct Zeros{
+    static inline __device__ T get() {
+        return T(0.f);
+    }
+};
+
+template<> 
+struct Zeros<float2>{
+    static inline __device__ float2 get() {
+        return make_float2(0.f, 0.f);
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename Elt_type, uint32_t NUM_ELT>
+struct Vec {
+
+    enum { BYTES = NUM_ELT * sizeof(Elt_type) };
+
+    using Vec_type = typename BytesToType<BYTES>::Type;
+
+    using Alias_type = union {
+        Vec_type vec;
+        Elt_type elt[NUM_ELT];
+    };
+
+    Alias_type data;
+
+    template<typename S>
+    inline __device__ void to(Vec<S, NUM_ELT> &other) {
+        #pragma unroll
+        for( int it = 0; it < NUM_ELT; it++ ) {
+            other.data.elt[it] = S(this->data.elt[it]);
+        }
+    }
+
+    template<typename Op>
+    inline __device__ void assign(const Op &op) {
+        #pragma unroll
+        for( int it = 0; it < NUM_ELT; it++ ) {
+            this->data.elt[it] = op(it);
+        }
+    }
+
+    inline __device__ void load_from(const void *base_ptr, const size_t idx) {
+        this->data.vec = static_cast<const Vec_type *>(base_ptr)[idx];
+    }
+
+    inline __device__ void store_to(void *base_ptr, const size_t idx) {
+        static_cast<Vec_type *>(base_ptr)[idx] = this->data.vec;
+    }
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<uint32_t CTAS_PER_ROW>
+struct InterCTASync {
+
+    template<typename Params>
+    inline __device__ InterCTASync(Params & params, uint32_t bidm, uint32_t bidn)
+        : phase_counter_(0)
+        , b0_(params.barrier + bidm) // The barrier for this group of CTAs.
+        , b1_(params.barrier + bidm + params.ctas_per_col) // The barrier for this group of CTAs.
+    {
+        // BARRIERS ARE ASSUMED TO BE INITIALIZED TO 0!
+    }
+
+    inline __device__ void spin_wait_(int *barrier, int step, int expected) {
+        asm volatile("red.release.gpu.global.add.s32 [%0], %1;" ::"l"(barrier), "r"(step));
+        for( int found = -1; found != expected; ) {
+            asm volatile("ld.global.acquire.gpu.b32 %0, [%1];" : "=r"(found) : "l"(barrier));
+        }
+    }
+
+    inline __device__ void sync(){
+        // ALL THREADS MUST ENTER!
+
+        // We switch barrier every iteration.
+        int *barrier = phase_counter_ & 0x1 ? b1_ : b0_;
+        // We decrement every other iteration.
+        bool dec = phase_counter_ & 0x2;
+        int step = dec ? -1 : 1;
+        int expected = dec ? 0 : CTAS_PER_ROW;
+        // There are only 4 phases: up/down for b0/b1.
+        phase_counter_ = (phase_counter_ + 1) & 0x3;
+
+        if( threadIdx.x == 0 ) {
+            spin_wait_(barrier, step, expected);
+        }
+        // CTA waits for thread 0
+        __syncthreads();
+    }
+
+    int phase_counter_;
+    int * b0_;
+    int * b1_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, uint32_t CTAS_PER_ROW, uint32_t WARPS_M, uint32_t WARPS_N>
+struct Reducer : public Reducer<T, 1, WARPS_M, WARPS_N> {
+
+    using InterCTASync = InterCTASync<CTAS_PER_ROW>;
+    using Base = Reducer<T, 1, WARPS_M, WARPS_N>;
+    using Type = typename Base::Type;
+
+    enum { SMEM_BYTES = Base::SMEM_BYTES };
+
+    enum { WS_BARRIER_BYTES = 2 * sizeof(int) };
+    enum { WS_DATA_BYTES = WARPS_M * CTAS_PER_ROW * sizeof(T) };
+
+    // size of the barriers + temporary result per CTA (multiply with CTAS_PER_ROW to get total)
+    enum { WORKSPACE_BYTES_PER_GROUP = Base::WORKSPACE_BYTES_PER_GROUP + WS_BARRIER_BYTES + WS_DATA_BYTES };
+
+    template<typename Params>
+    inline __device__ Reducer(Params & params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n, uint32_t lane, void * smem)
+        : Base(params, bidm, bidn, warp_m, warp_n, lane, smem) 
+        , inter_cta_(params, bidm, bidn)
+        , bidn_(bidn) // CTA id within the group.
+        , w0_(static_cast<T*>(params.workspace) + (bidm * WARPS_M + warp_m) * CTAS_PER_ROW)
+        , w1_(w0_ + params.ctas_per_col * WARPS_M * CTAS_PER_ROW)
+    {
+    }
+
+    template<typename Op>
+    inline __device__ T allreduce(T data, Op &op) {
+        data = Base::reduce(data, op);
+        // We switch workspace every iteration.
+        T *workspace = inter_cta_.phase_counter_ & 0x1 ? w1_ : w0_;
+
+        // Warp leaders 0 hold the CTA-local results.
+        if( this->warp_n_ == 0 && this->lane_ == 0 ) {
+            workspace[bidn_] = data;
+        }
+        inter_cta_.sync();
+        static_assert(CTAS_PER_ROW <= 32);
+        T total = Zeros<T>::get();
+        if(this->lane_ < CTAS_PER_ROW){
+            total = workspace[this->lane_];
+        }
+        total = Reducer<T, 1, 1, 1>::allreduce_(total, op);
+
+        return total;
+    }
+
+    InterCTASync inter_cta_;
+
+    T *w0_;
+    T *w1_;
+    int bidn_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, uint32_t WARPS_M>
+struct Reducer<T, 1, WARPS_M, 1> {
+
+    using Type = T;
+    enum { SMEM_BYTES = 0 };
+    enum { WORKSPACE_BYTES_PER_GROUP = 0 };
+
+    enum { THREADS_PER_WARP = 32 };
+
+    template<typename Params>
+    inline __device__ Reducer(Params & params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n, uint32_t lane, void * smem) 
+        : warp_n_(warp_n)
+        , lane_(lane)
+    {
+    }
+
+    template<typename Op>
+    static inline __device__ T allreduce_(T data, Op &op) {
+        #pragma unroll
+        for( int it = 1; it < THREADS_PER_WARP; it *= 2 ) {
+            data = op(data, warp_shuffle_xor(data, it));
+        }
+        return data;
+    }
+
+    template<typename Op>
+    inline __device__ T allreduce(T data, Op &op) {
+        return allreduce_(data, op);
+    }
+
+    template<typename Op>
+    inline __device__ T reduce(T data, Op &op){
+        // only lane 0 holds the result!
+        #pragma unroll
+        for( int it = THREADS_PER_WARP / 2; it > 0; it /= 2 ) {
+            data = op(data, warp_shuffle_down(data, it));
+        }  
+        return data;
+    }
+    int warp_n_;
+    int lane_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, uint32_t WARPS_M, uint32_t WARPS_N>
+struct Reducer<T, 1, WARPS_M, WARPS_N> : public Reducer<T, 1, WARPS_M, 1> {
+
+    using Base = Reducer<T, 1, WARPS_M, 1>;
+
+    using Type = T;
+
+    enum { SMEM_BYTES = Base::SMEM_BYTES + WARPS_M * WARPS_N * sizeof(T) * 2 };
+    enum { WORKSPACE_BYTES_PER_GROUP = 0 };
+
+    enum { THREADS_PER_WARP = 32 };
+
+    template<typename Params>
+    inline __device__ Reducer(Params & params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n, uint32_t lane, void * smem) 
+        : Base(params, bidm, bidn, warp_m, warp_n, lane, smem) 
+        , use0_(true)
+    {
+        smem0_ = &static_cast<T *>(smem)[warp_m * WARPS_N];
+        smem1_ = smem0_ + WARPS_M * WARPS_N;
+    }
+
+    template<typename Op>
+    inline __device__ T allreduce(T data, Op & op) {
+        T * smem = use0_ ? smem0_ : smem1_;
+        use0_ = !use0_;
+        data = Base::reduce(data, op);
+        if( this->lane_ == 0 ) {
+            smem[this->warp_n_] = data;
+        }
+        __syncthreads();
+        T out = Zeros<T>::get();
+        #pragma unroll
+        for( int it = 0; it < WARPS_N; it++ ) {
+            out = op(out, smem[it]);
+        }
+        return out;
+    }
+
+    template<typename Op>
+    inline __device__ T reduce(T data, Op &op) {
+        T * smem = use0_ ? smem0_ : smem1_;
+        use0_ = !use0_;
+        // only intra-CTA group leader holds the result!
+        data = Base::reduce(data, op);
+        if( this->lane_ == 0 ) {
+            smem[this->warp_n_] = data;
+        }
+        __syncthreads();
+        T out = Zeros<T>::get();
+        if( this->warp_n_ == 0 && this->lane_ == 0 ) {
+            #pragma unroll
+            for( int it = 0; it < WARPS_N; it++ ) {
+                out = op(out, smem[it]);
+            }
+        }
+        return out;
+    }
+
+    T * smem0_;
+    T * smem1_;
+    bool use0_;
+
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+ 
+template<typename T>
+inline __device__ void warp_chan_upd_dynamic(T &m_a, T &m2_a, T &n_a, int num_active){
+    //Assume at least leftmost is valid and init: step = next_pow2(num_active) / 2 (might get NaN otherwise)
+    int highest_bit_set = (8 * sizeof(num_active)) - __clz(num_active - 1);
+    
+    #pragma unroll
+    for( int step = (1 << (highest_bit_set - 1)); step > 0; step /= 2 ) {
+        // Exchange
+        T n_b = warp_shuffle_down(n_a, step);
+        T m_b = warp_shuffle_down(m_a, step);
+        T m2_b = warp_shuffle_down(m2_a, step);
+
+        // Update
+        const T n_ab = n_a + n_b; // We can handle one of them being 0, not both.
+        const T rn_ab = 1.f / n_ab; // Might have different n per thread, otherwise this would simplify :(
+        const T delta = m_a - m_b;
+        const float m2_ab = m2_a + m2_b + delta * delta * n_a * n_b * rn_ab;
+        const float m_ab = (n_a * m_a + n_b * m_b) * rn_ab;
+
+        n_a = n_ab;
+        m_a = m_ab;
+        m2_a = m2_ab;
+    }
+    // Intra-warp broadcast (only lane 0 has valid stats).
+    m_a = __shfl_sync(uint32_t(-1), m_a, 0);
+    m2_a = __shfl_sync(uint32_t(-1), m2_a, 0);
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, uint32_t CTAS_PER_ROW, uint32_t WARPS_M, uint32_t WARPS_N>
+struct Stats {
+    // This could be done generically with the Reducer. But then we would have to exchange 3 instead of 2 fields.
+
+    using InterCTASync = InterCTASync<CTAS_PER_ROW>;
+    using BlockStats = Stats<T, 1, WARPS_M, WARPS_N>;
+    using stats_t = typename BlockStats::stats_t;
+
+    enum { SMEM_BYTES = BlockStats::SMEM_BYTES };
+
+    template<typename Params>
+    inline __device__ Stats(Params & params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n, uint32_t lane, void * smem) 
+        : inter_cta_(params, bidm, bidn)
+        , block_stats_(params, bidm, bidn, warp_m, warp_n, lane, smem)
+        , bidn_(bidn) // CTA id within the group.
+        , w0_(static_cast<stats_t*>(params.workspace) + (bidm * WARPS_M + warp_m) * CTAS_PER_ROW)
+        , w1_(w0_ + params.ctas_per_col * WARPS_M * CTAS_PER_ROW)
+        , warp_n_(warp_n)
+        , lane_(lane)
+    {
+    }
+
+    template<uint32_t N>
+    inline __device__ stats_t compute(const T (&elts)[N], const T rn) {
+        constexpr T ELTS_PER_ROW_PER_CTA = N * WARPS_N * THREADS_PER_WARP;
+        // TODO rn is not really needed here..
+        constexpr T block_rn = 1.f / T(ELTS_PER_ROW_PER_CTA);
+        stats_t block_stats = block_stats_.compute(elts, block_rn);
+
+        stats_t *workspace = inter_cta_.phase_counter_ & 0x1 ? w1_ : w0_;
+
+        if( warp_n_ == 0 && lane_ == 0 ) {
+            workspace[bidn_] = block_stats;
+        }
+
+        // Wait for all CTAS_PER_ROW CTAS in the group to have written their result.
+        inter_cta_.sync();
+
+        T n = Zeros<T>::get();
+        T m = Zeros<T>::get();
+        T m2 = Zeros<T>::get();
+
+        // Assume CTA group size in N less than 32, such that we can finalize with a single warp.
+        static_assert(CTAS_PER_ROW <= 32);
+
+        // Every warp does the final reduction locally. 
+        if( lane_ < CTAS_PER_ROW ) {
+            stats_t result = workspace[lane_];
+            n = ELTS_PER_ROW_PER_CTA;
+            m = layer_norm::Get<0>::of<stats_t, T>(result);
+            m2 = layer_norm::Get<1>::of<stats_t, T>(result);
+        }
+
+        warp_chan_upd_dynamic(m, m2, n, CTAS_PER_ROW);
+
+        return { m, m2 };
+    }
+
+    InterCTASync inter_cta_;
+    BlockStats block_stats_;
+
+    stats_t *w0_;
+    stats_t *w1_;
+    int bidn_;
+    int warp_n_;
+    int lane_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, uint32_t WARPS_M, uint32_t WARPS_N>
+struct Stats<T, 1, WARPS_M, WARPS_N> {
+
+    using WarpStats = Stats<T, 1, WARPS_M, 1>;
+    using stats_t = typename WarpStats::stats_t;
+
+    enum { SMEM_BYTES = WARPS_M * WARPS_N * sizeof(stats_t) * 2 };
+
+    template<typename Params>
+    inline __device__ Stats(Params & params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n, uint32_t lane, void * smem) 
+        : warp_stats_(params, bidm, bidn, warp_m, warp_n, lane, smem)
+        , use0_(true)
+    {
+        smem0_ = static_cast<stats_t*>(smem) + warp_m * WARPS_N;
+        smem1_ = smem0_ + WARPS_M * WARPS_N;
+    }
+
+    template<uint32_t N>
+    inline __device__ stats_t compute(const T (&elts)[N], const T rn) {
+        stats_t * smem = use0_ ? smem0_ : smem1_;
+        use0_ = !use0_;
+        // Compute warp local for all WARPS_N
+        constexpr T warp_rn = 1.f / T(N * THREADS_PER_WARP);
+        stats_t warp_stats = warp_stats_.compute(elts, warp_rn);
+
+        //Each warp warp leader stores its stats
+        const auto warp_n = warp_stats_.reducer_.warp_n_;
+        const auto lane = warp_stats_.reducer_.lane_;
+        if( lane == 0 ) {
+            smem[warp_n] = warp_stats;
+        }
+        __syncthreads();
+
+        T n = Zeros<T>::get();
+        T m = Zeros<T>::get();
+        T m2 = Zeros<T>::get();
+
+        // Assume that there are less than 32 warps, such that we can finalize with a single warp
+        static_assert(WARPS_N <= 32);
+        if(lane < WARPS_N){
+            stats_t result = smem[lane];
+            n = N * THREADS_PER_WARP;
+            m = layer_norm::Get<0>::of<stats_t, T>(result);
+            m2 = layer_norm::Get<1>::of<stats_t, T>(result);
+        }
+
+        warp_chan_upd_dynamic(m, m2, n, WARPS_N);
+
+        return { m, m2 };
+    }
+    WarpStats warp_stats_;
+    stats_t * smem0_;
+    stats_t * smem1_;
+    bool use0_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template<typename T, uint32_t WARPS_M>
+struct Stats<T, 1, WARPS_M, 1> {
+
+    using stats_t = typename TypeToVec2<T>::Type;
+    // The simple Warp reducer.
+    using Reducer = Reducer<T, 1, WARPS_M, 1>;
+
+    enum { SMEM_BYTES = 0 };
+
+    template<typename Params>
+    inline __device__ Stats(Params & params, uint32_t bidm, uint32_t bidn, uint32_t warp_m, uint32_t warp_n, uint32_t lane, void * smem) 
+        : reducer_(params, bidm, bidn, warp_m, warp_n, lane, smem)
+    {
+    }
+
+    template<uint32_t N>
+    inline __device__ stats_t compute(const T (&elts)[N], const T rn) {
+
+        auto sum = Sum<T>();
+
+        T m = Zeros<T>::get();
+        #pragma unroll
+        for( int it = 0; it < N; it++ ) {
+            m += elts[it];
+        }
+        m = reducer_.allreduce(m, sum) * rn;
+
+        T m2 = Zeros<T>::get();
+        #pragma unroll
+        for( int it = 0; it < N; it++ ) {
+            T diff = (elts[it] - m);
+            m2 += diff * diff;
+        }
+        m2 = reducer_.allreduce(m2, sum);
+
+        return {m, m2};
+    }
+
+    Reducer reducer_;
+};
+
+////////////////////////////////////////////////////////////////////////////////////////////////////
+
+}  // namespace layer_norm
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout.cpp
new file mode 100644
index 0000000000..d2573bdf1a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout.cpp
@@ -0,0 +1,91 @@
+#include <torch/extension.h>
+#include <cuda_fp16.h>
+#include <vector>
+
+namespace multihead_attn {
+namespace fused_softmax {
+namespace additive_mask_softmax_dropout {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& input, 
+                               const half*       pad_mask,
+                               float                dropout_prob
+                                                  );
+
+torch::Tensor bwd_cuda(
+		               int heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  );
+
+// C++ interface
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> fwd(
+ 			       bool 				use_mask,
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& input,
+                               torch::Tensor const& pad_mask,
+                               float                dropout_prob
+                                                 )
+{
+  AT_ASSERTM(input.dim()         == 3, "expected 3D tensor");
+  AT_ASSERTM(input.type().scalarType()         == at::ScalarType::Half, "Only HALF is supported");
+
+  if (use_mask) {
+  	AT_ASSERTM(pad_mask.dim()                     == 2,                    "expected 2D tensor");
+  	AT_ASSERTM(pad_mask.type().scalarType()       == at::ScalarType::Half, "Only BYTE is supported");
+  }
+
+  return fwd_cuda(
+                                 is_training,
+                                 heads, 
+                                 input, 
+                                 use_mask ? static_cast<const half*>(pad_mask.data_ptr()) : nullptr, 
+                                 dropout_prob
+                                );
+}
+
+torch::Tensor bwd(
+		               bool use_mask,
+		               int heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  )
+{
+  AT_ASSERTM(output_grads.dim()      == 3, "expected 3D tensor");
+  AT_ASSERTM(softmax_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(dropout_mask.dim()      == 3, "expected 3D tensor");
+
+  AT_ASSERTM(output_grads.type().scalarType()      == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(softmax_results.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+//  AT_ASSERTM(dropout_mask.type().scalarType()      == at::ScalarType::Byte, "Only BYTE is supported");
+
+  return bwd_cuda(
+		                 heads,
+                                 output_grads,
+                                 softmax_results, 
+                                 dropout_mask, 
+                                 dropout_prob
+                                );
+}
+
+} // end namespace mask_softmax_dropout
+} // end namespace fused_softmax
+} // end namespace multihead_attn
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &multihead_attn::fused_softmax::additive_mask_softmax_dropout::fwd, "Self Multihead Attention masked softmax dropout -- Forward.");
+  m.def("backward", &multihead_attn::fused_softmax::additive_mask_softmax_dropout::bwd, "Self Multihead Attention masked softmax dropout -- Backward.");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout_cuda.cu
new file mode 100644
index 0000000000..05645b1653
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout_cuda.cu
@@ -0,0 +1,131 @@
+#include <vector>
+#include <math.h>
+#include <iostream>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include "softmax.h"
+#include "dropout.h"
+
+// symbol to be automatically resolved by PyTorch libs
+extern THCState *state;
+
+namespace multihead_attn {
+namespace fused_softmax {
+namespace additive_mask_softmax_dropout {
+
+std::vector<torch::Tensor> fwd_cuda(
+			       bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& input, 
+                               const half*       pad_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   attn_batches   = input.size(0);
+  const int   sequences      = attn_batches / heads;
+  const int   q_seq_len      = input.size(1);
+  const int   k_seq_len      = q_seq_len;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+
+  // There is no reason to use more than one stream as every kernel is 
+  // sequentially dependent
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // 3 Intermediate Results + Output (Note: dropout intermediates are generated by ATen library code)
+  auto act_options  = input.options().requires_grad(false);
+  auto mask_options = act_options.dtype(torch::kUInt8);
+
+  torch::Tensor softmax_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_mask      = torch::empty({attn_batches, q_seq_len, k_seq_len},   mask_options);
+
+  // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)
+  void* input_ptr = static_cast<void*>(input.data_ptr());
+  void* softmax_results_ptr = static_cast<void*>(softmax_results.data_ptr());
+
+  // Padded Softmax
+  bool softmax_success = false;
+  if (pad_mask == nullptr) {
+    softmax_success = dispatch_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(input_ptr),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  } else {
+      softmax_success = dispatch_additive_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(input_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             attn_batches*q_seq_len/sequences);
+  }
+
+
+  if (is_training) {
+    //use at:: function so that C++ version generates the same random mask as python version
+    auto dropout_tuple = at::_fused_dropout(softmax_results, 1.0f-dropout_prob);
+    dropout_results = std::get<0>(dropout_tuple);
+    dropout_mask = std::get<1>(dropout_tuple);
+  }
+
+  // Matmul2
+
+  return {
+           dropout_results,  
+           dropout_mask, 
+           softmax_results
+         };
+}
+
+torch::Tensor bwd_cuda(
+		               int heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& softmax_results, 
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   attn_batches   = output_grads.size(0);
+  const int   q_seq_len      = output_grads.size(1);
+  const int   k_seq_len      = q_seq_len;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+  // TODO: Streams can be used in Backprop but I haven't added more than one
+  // in my first attempt to create the code
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // Output Tensor Allocations
+//  torch::Tensor input_grads         = torch::empty_like(output_grads);
+
+  // Apply Dropout Mask and Scale by Dropout Probability 
+  // Softmax Grad
+  dispatch_masked_scale_softmax_backward_stream<half, half, float,false>(
+                             static_cast<half*>(output_grads.data_ptr()), 
+                             static_cast<half*>(output_grads.data_ptr()), 
+                             reinterpret_cast<half const*>(softmax_results.data_ptr()),
+			     static_cast<uint8_t const*>(dropout_mask.data_ptr()),
+			     1.0/(1.0-dropout_prob),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len, stream);
+//backward pass is completely in-place
+  return output_grads;
+}
+}
+}
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/dropout.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/dropout.h
new file mode 100644
index 0000000000..f10a50f79a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/dropout.h
@@ -0,0 +1,306 @@
+#include <ATen/ATen.h>
+
+#ifdef OLD_GENERATOR
+#include <ATen/CUDAGenerator.h>
+#else
+#include <ATen/CUDAGeneratorImpl.h>
+#endif
+
+#include <ATen/cuda/CUDAContext.h>
+#include <curand_kernel.h>
+
+const int UNROLL = 4;
+
+template <
+          typename scalar_t,
+          typename accscalar_t,
+          typename IndexType
+         >
+__global__ void apex_fused_dropout_kernel(scalar_t const                *inputs,
+                                          scalar_t                      *outputs,
+                                          uint8_t                       *mask,
+                                          IndexType                      totalElements, 
+		                                  accscalar_t                    p, 
+		                                  std::pair<uint64_t, uint64_t>  seeds
+                                         ) 
+{
+  accscalar_t pinv = accscalar_t(1)/p;
+  IndexType idx = blockIdx.x * blockDim.x + threadIdx.x;
+
+  curandStatePhilox4_32_10_t state;
+  curand_init(
+      seeds.first,
+      idx,
+      seeds.second,
+      &state);
+
+  IndexType rounded_size = ((totalElements - 1)/(blockDim.x * gridDim.x * UNROLL)+1) * blockDim.x * gridDim.x * UNROLL;
+  for (IndexType linearIndex = idx;
+       linearIndex < rounded_size;
+       linearIndex += gridDim.x * blockDim.x*UNROLL) {
+       float4 rand = curand_uniform4(&state);
+       scalar_t src[UNROLL];
+       rand.x = rand.x <= p;
+       rand.y = rand.y <= p;
+       rand.z = rand.z <= p;
+       rand.w = rand.w <= p;
+
+       for (int ii = 0; ii < UNROLL; ii++) {
+           IndexType li = linearIndex + blockDim.x * gridDim.x * ii;
+           if (li < totalElements) {
+               src[ii] = inputs[li];
+           }
+       }
+       for (int ii = 0; ii < UNROLL; ii++) {
+           IndexType li = linearIndex + blockDim.x * gridDim.x * ii;
+           if (li < totalElements) {
+	           outputs[li] = src[ii]*(&rand.x)[ii]*pinv;
+               mask[li]    = (uint8_t)(&rand.x)[ii];
+           }
+       }
+       __syncthreads();
+  }
+}
+
+template <
+          typename scalar_t,
+          typename accscalar_t,
+          typename IndexType
+         >
+__global__ void apex_dropout_add_kernel(scalar_t const                *inputs,
+                                        scalar_t const                *add_inputs,
+                                        scalar_t                      *outputs,
+                                        uint8_t                       *mask,
+                                        IndexType                      totalElements, 
+		                                accscalar_t                    p, 
+		                                std::pair<uint64_t, uint64_t>  seeds
+                                       ) 
+{
+  accscalar_t pinv = accscalar_t(1)/p;
+  IndexType idx = blockIdx.x * blockDim.x + threadIdx.x;
+
+  curandStatePhilox4_32_10_t state;
+  curand_init(
+      seeds.first,
+      idx,
+      seeds.second,
+      &state);
+
+  IndexType rounded_size = ((totalElements - 1)/(blockDim.x * gridDim.x * UNROLL)+1) * blockDim.x * gridDim.x * UNROLL;
+  for (IndexType linearIndex = idx;
+       linearIndex < rounded_size;
+       linearIndex += gridDim.x * blockDim.x*UNROLL) {
+       float4 rand = curand_uniform4(&state);
+       scalar_t src[UNROLL];
+       scalar_t add_src[UNROLL];
+       rand.x = rand.x <= p;
+       rand.y = rand.y <= p;
+       rand.z = rand.z <= p;
+       rand.w = rand.w <= p;
+       for (int ii = 0; ii < UNROLL; ii++) {
+           IndexType li = linearIndex + blockDim.x * gridDim.x * ii;
+           if (li < totalElements) {
+               src[ii]     = inputs[li];
+               add_src[ii] = add_inputs[li];
+           }
+       }
+       for (int ii = 0; ii < UNROLL; ii++) {
+           IndexType li = linearIndex + blockDim.x * gridDim.x * ii;
+           if (li < totalElements) {
+	           accscalar_t int1 = src[ii] * (&rand.x)[ii] * pinv;
+	           outputs[li] = static_cast<scalar_t>(static_cast<accscalar_t>(add_src[ii]) + int1);
+               mask[li]    = (uint8_t)(&rand.x)[ii];
+           }
+       }
+       __syncthreads();
+  }
+}
+
+template <
+          typename scalar_t,
+          typename accscalar_t,
+          typename IndexType
+         >
+__global__ void apex_add_kernel(          scalar_t const                *inputs,
+                                        scalar_t const                *add_inputs,
+                                        scalar_t                      *outputs,
+                                        IndexType                      totalElements
+                             ) 
+{
+  IndexType idx = blockIdx.x * blockDim.x + threadIdx.x;
+  IndexType rounded_size = ((totalElements - 1)/(blockDim.x * gridDim.x * UNROLL)+1) * blockDim.x * gridDim.x * UNROLL;
+  for (IndexType linearIndex = idx;
+       linearIndex < rounded_size;
+       linearIndex += gridDim.x * blockDim.x*UNROLL) {
+       scalar_t src[UNROLL];
+       scalar_t add_src[UNROLL];
+       for (int ii = 0; ii < UNROLL; ii++) {
+           IndexType li = linearIndex + blockDim.x * gridDim.x * ii;
+           if (li < totalElements) {
+               src[ii]     = inputs[li];
+               add_src[ii] = add_inputs[li];
+           }
+       }
+       for (int ii = 0; ii < UNROLL; ii++) {
+           IndexType li = linearIndex + blockDim.x * gridDim.x * ii;
+           if (li < totalElements) {
+	           outputs[li] = src[ii] + add_src[ii];
+           }
+       }
+       __syncthreads();
+  }
+}
+
+template<typename scalar_t, 
+		 typename accscalar_t, 
+		 typename IndexType
+		>
+__global__ void apex_masked_scale_kernel(scalar_t const *inputs, 
+                                         scalar_t       *outputs, 
+                                         uint8_t const  *mask, 
+                                         IndexType       totalElements,
+                                         accscalar_t     scale
+                                        )
+{
+  IndexType idx          = blockIdx.x * blockDim.x + threadIdx.x;
+  IndexType rounded_size = ((totalElements - 1)/(blockDim.x * gridDim.x * UNROLL)+1) * blockDim.x * gridDim.x * UNROLL;
+  for (IndexType linearIndex = idx;
+       linearIndex < rounded_size;
+       linearIndex += gridDim.x * blockDim.x*UNROLL) 
+  {
+       scalar_t src[UNROLL];
+       scalar_t msk[UNROLL];
+       for (int ii = 0; ii < UNROLL; ii++) {
+           IndexType li = linearIndex + blockDim.x * gridDim.x * ii;
+           if (li < totalElements) {
+               src[ii] = static_cast<scalar_t>(inputs[li]);
+               msk[ii] = static_cast<scalar_t>(mask[li]);
+           }
+       }
+       for (int ii = 0; ii < UNROLL; ii++) {
+           IndexType li = linearIndex + blockDim.x * gridDim.x * ii;
+           if (li < totalElements) {
+               outputs[li] = static_cast<accscalar_t>(src[ii]) * scale * static_cast<accscalar_t>(msk[ii]);
+           }
+       }
+  }
+}
+
+template <
+          typename scalar_t,
+          typename accscalar_t,
+          typename IndexType
+         >
+void apex_fused_dropout_cuda(scalar_t const *inputs,
+                           scalar_t       *outputs,
+                           uint8_t        *mask,
+                           IndexType       totalElements, 
+		                   accscalar_t     p)
+{
+  auto gen = at::cuda::detail::getDefaultCUDAGenerator();
+  
+  int block_size = 256;
+  dim3 dim_block(block_size);
+  dim3 grid((totalElements + block_size -1)/block_size);
+  unsigned int blocks_per_sm = at::cuda::getCurrentDeviceProperties()->maxThreadsPerMultiProcessor/block_size;
+  grid.x = std::min((unsigned int)at::cuda::getCurrentDeviceProperties()->multiProcessorCount * blocks_per_sm, grid.x);
+
+  //number of times random will be generated per thread, to offset philox counter in the random state
+  int64_t counter_offset = ((totalElements - 1)/(block_size*grid.x*UNROLL)+1)*UNROLL;
+  std::pair<uint64_t, uint64_t> rng_engine_inputs;
+  {
+    // See Note [Acquire lock when using random generators]
+#ifdef OLD_GENERATOR
+    std::lock_guard<std::mutex> lock(gen->mutex_);
+    rng_engine_inputs = gen->philox_engine_inputs(counter_offset);
+#else
+    std::lock_guard<std::mutex> lock(gen.mutex());
+    rng_engine_inputs = at::check_generator<at::CUDAGeneratorImpl>(gen)->philox_engine_inputs(counter_offset);
+#endif
+  }
+
+  apex_fused_dropout_kernel<scalar_t, accscalar_t, IndexType><<<grid, dim_block, 0, at::cuda::getCurrentCUDAStream()>>>(inputs, outputs, mask, totalElements, p, rng_engine_inputs);
+  C10_CUDA_CHECK(cudaGetLastError());
+}
+
+template <
+          typename scalar_t,
+          typename accscalar_t,
+          typename IndexType
+         >
+void apex_dropout_add_cuda(scalar_t const *inputs,
+                           scalar_t const *add_inputs,
+                           scalar_t       *outputs,
+                           uint8_t        *mask,
+                           IndexType       totalElements, 
+		                   accscalar_t     p)
+{
+  auto gen = at::cuda::detail::getDefaultCUDAGenerator();
+  
+  int block_size = 256;
+  dim3 dim_block(block_size);
+  dim3 grid((totalElements + block_size -1)/block_size);
+  unsigned int blocks_per_sm = at::cuda::getCurrentDeviceProperties()->maxThreadsPerMultiProcessor/block_size;
+  grid.x = std::min((unsigned int)at::cuda::getCurrentDeviceProperties()->multiProcessorCount * blocks_per_sm, grid.x);
+
+  //number of times random will be generated per thread, to offset philox counter in the random state
+  int64_t counter_offset = ((totalElements - 1)/(block_size*grid.x*UNROLL)+1)*UNROLL;
+  std::pair<uint64_t, uint64_t> rng_engine_inputs;
+  {
+    // See Note [Acquire lock when using random generators]
+#ifdef OLD_GENERATOR
+    std::lock_guard<std::mutex> lock(gen->mutex_);
+    rng_engine_inputs = gen->philox_engine_inputs(counter_offset);
+#else
+    std::lock_guard<std::mutex> lock(gen.mutex());
+    rng_engine_inputs = at::check_generator<at::CUDAGeneratorImpl>(gen)->philox_engine_inputs(counter_offset);
+#endif
+  }
+
+  apex_dropout_add_kernel<scalar_t, accscalar_t, IndexType><<<grid, dim_block, 0, at::cuda::getCurrentCUDAStream()>>>(inputs, add_inputs, outputs, mask, totalElements, p, rng_engine_inputs);
+  C10_CUDA_CHECK(cudaGetLastError());
+}
+
+template <
+          typename scalar_t,
+          typename accscalar_t,
+          typename IndexType
+         >
+void apex_add_cuda(scalar_t const *inputs,
+                   scalar_t const *add_inputs,
+                   scalar_t       *outputs,
+                   IndexType       totalElements
+		          )
+{
+  int block_size = 256;
+  dim3 dim_block(block_size);
+  dim3 grid((totalElements + block_size -1)/block_size);
+  unsigned int blocks_per_sm = at::cuda::getCurrentDeviceProperties()->maxThreadsPerMultiProcessor/block_size;
+  grid.x = std::min((unsigned int)at::cuda::getCurrentDeviceProperties()->multiProcessorCount * blocks_per_sm, grid.x);
+
+  apex_add_kernel<scalar_t, accscalar_t, IndexType><<<grid, dim_block, 0, at::cuda::getCurrentCUDAStream()>>>(inputs, add_inputs, outputs, totalElements);
+  C10_CUDA_CHECK(cudaGetLastError());
+}
+
+template<typename scalar_t, 
+         typename accscalar_t, 
+         typename IndexType
+        >
+void apex_masked_scale_cuda(scalar_t const *inputs, 
+                          scalar_t       *outputs, 
+                          uint8_t const  *mask, 
+                          IndexType       totalElements,
+                          accscalar_t     scale
+                         )
+{
+  int block_size = 256;
+  dim3 dim_block(block_size);
+  dim3 grid((totalElements + block_size -1)/block_size);
+  unsigned int blocks_per_sm = at::cuda::getCurrentDeviceProperties()->maxThreadsPerMultiProcessor/block_size;
+  grid.x = std::min((unsigned int)at::cuda::getCurrentDeviceProperties()->multiProcessorCount * blocks_per_sm, grid.x);
+
+  apex_masked_scale_kernel<scalar_t, accscalar_t, IndexType><<<grid, dim_block, 0, at::cuda::getCurrentCUDAStream()>>>(inputs, outputs, mask, totalElements, scale);
+  C10_CUDA_CHECK(cudaGetLastError());
+}
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn.cpp
new file mode 100644
index 0000000000..35d4d11095
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn.cpp
@@ -0,0 +1,156 @@
+#include <torch/extension.h>
+#include <vector>
+
+namespace multihead_attn {
+namespace encdec {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,  
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                                  );
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_q_results,
+                               torch::Tensor const& input_lin_kv_results,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  );
+
+// C++ interface
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> fwd(
+                               bool                 use_mask,
+                               bool                 use_time_mask,
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& pad_mask,
+                               float                dropout_prob
+                                                 )
+{
+  AT_ASSERTM(inputs_q.dim()         == 3, "expected 3D tensor");
+  AT_ASSERTM(inputs_kv.dim()        == 3, "expected 3D tensor");
+  AT_ASSERTM(input_weights_q.dim()  == 2, "expected 2D tensor");
+  AT_ASSERTM(input_weights_kv.dim() == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim()   == 2, "expected 2D tensor");
+
+  AT_ASSERTM(inputs_q.type().scalarType()         == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(inputs_kv.type().scalarType()        == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights_q.type().scalarType()  == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights_kv.type().scalarType() == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+  
+  if (use_mask) {
+  	AT_ASSERTM(pad_mask.dim()                     == 2,                    "expected 2D tensor");
+  	AT_ASSERTM(pad_mask.type().scalarType()       == at::ScalarType::Byte, "Only BYTE is supported");
+  }
+  
+  return fwd_cuda(
+                                 use_time_mask,
+                                 is_training,
+                                 heads, 
+                                 inputs_q, 
+                                 inputs_kv, 
+                                 input_weights_q, 
+                                 input_weights_kv, 
+                                 output_weights, 
+                                 use_mask ? static_cast<const uint8_t*>(pad_mask.data_ptr()) : nullptr, 
+                                 dropout_prob
+                                );
+}
+
+std::vector<torch::Tensor> bwd(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_q_results,
+                               torch::Tensor const& input_lin_kv_results,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  )
+{
+  AT_ASSERTM(output_grads.dim()         == 3, "expected 3D tensor");
+  AT_ASSERTM(matmul2_results.dim()      == 3, "expected 3D tensor");
+  AT_ASSERTM(dropout_results.dim()      == 3, "expected 3D tensor");
+  AT_ASSERTM(softmax_results.dim()      == 3, "expected 3D tensor");
+  AT_ASSERTM(input_lin_q_results.dim()  == 3, "expected 3D tensor");
+  AT_ASSERTM(input_lin_kv_results.dim() == 3, "expected 3D tensor");
+  AT_ASSERTM(inputs_q.dim()             == 3, "expected 3D tensor");
+  AT_ASSERTM(inputs_kv.dim()            == 3, "expected 3D tensor");
+  AT_ASSERTM(input_weights_q.dim()      == 2, "expected 2D tensor");
+  AT_ASSERTM(input_weights_kv.dim()     == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim()       == 2, "expected 2D tensor");
+  AT_ASSERTM(dropout_mask.dim()         == 3, "expected 3D tensor");
+  
+  AT_ASSERTM(output_grads.type().scalarType()         == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(matmul2_results.type().scalarType()      == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(dropout_results.type().scalarType()      == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(softmax_results.type().scalarType()      == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_lin_q_results.type().scalarType()  == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_lin_kv_results.type().scalarType() == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(inputs_q.type().scalarType()             == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(inputs_kv.type().scalarType()            == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights_q.type().scalarType()      == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights_kv.type().scalarType()     == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType()       == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(dropout_mask.type().scalarType()         == at::ScalarType::Byte, "Only BYTE is supported");
+  
+  return bwd_cuda(
+                                 heads, 
+                                 output_grads,
+                                 matmul2_results,
+                                 dropout_results,
+                                 softmax_results, 
+                                 input_lin_q_results, 
+                                 input_lin_kv_results, 
+                                 inputs_q, 
+                                 inputs_kv, 
+                                 input_weights_q,
+                                 input_weights_kv,
+                                 output_weights,
+                                 dropout_mask, 
+                                 dropout_prob
+                                );
+}
+
+} // end namespace cublas_gemmex
+} // end namespace encdec 
+} // end namespace multihead_attn
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &multihead_attn::encdec::cublas_gemmex::fwd, "Encdec Multihead Attention Forward.");
+  m.def("backward", &multihead_attn::encdec::cublas_gemmex::bwd, "Encdec Multihead Attention Backward.");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_cuda.cu
new file mode 100644
index 0000000000..5e5e23d85f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_cuda.cu
@@ -0,0 +1,556 @@
+#include <vector>
+#include <math.h>
+#include <iostream>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include "strided_batched_gemm.h"
+#include "softmax.h"
+#include "dropout.h"
+#include "layer_norm.h"
+
+// symbol to be automatically resolved by PyTorch libs
+extern THCState *state;
+
+namespace multihead_attn {
+namespace encdec {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,
+							   bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim         = inputs_q.size(2);
+  const int   sequences         = inputs_q.size(1);
+  const int   q_seq_len         = inputs_q.size(0);
+  const int   k_seq_len         = inputs_kv.size(0);
+  const int   batches_q         = sequences * q_seq_len;
+  const int   batches_kv        = sequences * k_seq_len;
+  const int   head_dim          = embed_dim / heads;
+  const int   output_lin_q_dim  = embed_dim;
+  const int   output_lin_kv_dim = 2 * embed_dim;
+  const int   attn_batches      = heads * sequences;
+  const int   lead_dim_q        = attn_batches * head_dim;
+  const int   lead_dim_kv       = attn_batches * 2 *head_dim;
+  const int   batch_stride_q    = head_dim;
+  const int   batch_stride_kv   = 2 * head_dim;
+  const int   dropout_elems     = attn_batches * q_seq_len * k_seq_len;
+  const float alpha             = 1.0;
+  const float beta              = 0.0;
+  const float scale             = 1.0 / sqrt(static_cast<float>(head_dim));
+ 
+  // There is no reason to use more than one stream as every kernel is 
+  // sequentially dependent
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // 3 Intermediate Results + Output (Note: dropout intermediates are generated by ATen library code)
+  auto act_options  = inputs_q.options().requires_grad(false);
+  auto mask_options = act_options.dtype(torch::kUInt8);
+
+  torch::Tensor input_lin_q_results  = torch::empty({q_seq_len, sequences, output_lin_q_dim},  act_options);
+  torch::Tensor input_lin_kv_results = torch::empty({k_seq_len, sequences, output_lin_kv_dim}, act_options);
+  torch::Tensor softmax_results      = torch::empty({attn_batches, q_seq_len, k_seq_len},      act_options);
+  torch::Tensor dropout_results      = torch::empty({attn_batches, q_seq_len, k_seq_len},      act_options);
+  torch::Tensor dropout_mask         = torch::empty({attn_batches, q_seq_len, k_seq_len},      mask_options);
+  torch::Tensor matmul2_results      = torch::empty({q_seq_len, attn_batches, head_dim},       act_options);
+  torch::Tensor outputs              = torch::empty_like(inputs_q, act_options);
+
+  // Input Linear Results Pointers to Q, K, and V of interviewed activations
+  void* q_lin_results_ptr   = static_cast<void*>(input_lin_q_results.data_ptr());
+  void* k_lin_results_ptr   = static_cast<void*>(input_lin_kv_results.data_ptr());
+  void* v_lin_results_ptr   = static_cast<void*>(static_cast<half*>(input_lin_kv_results.data_ptr()) + head_dim);
+
+  // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)
+  void* softmax_results_ptr = static_cast<void*>(softmax_results.data_ptr());
+  
+  char a_layout_t{'t'};
+  char a_layout_n{'n'};
+  char b_layout_n{'n'};
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+  // Input Linear Q Fwd
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             output_lin_q_dim, 
+                             batches_q, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights_q.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(inputs_q.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             q_lin_results_ptr,
+                             CUDA_R_16F, 
+                             output_lin_q_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // Input Linear KV Fwd
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             output_lin_kv_dim, 
+                             batches_kv, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights_kv.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(inputs_kv.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             k_lin_results_ptr,
+                             CUDA_R_16F, 
+                             output_lin_kv_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             scale, 
+                             static_cast<const half*>(k_lin_results_ptr), 
+                             lead_dim_kv, 
+                             batch_stride_kv, 
+                             static_cast<const half*>(q_lin_results_ptr),
+                             lead_dim_q, 
+                             batch_stride_q, 
+                             beta, 
+                             static_cast<half*>(softmax_results_ptr), 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             attn_batches);
+
+  // Padded Softmax
+  bool softmax_success = false;
+  if (pad_mask == nullptr) {
+    softmax_success = dispatch_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  } else {
+    if (use_time_mask) {
+      softmax_success = dispatch_time_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             q_seq_len);
+    } else {
+      softmax_success = dispatch_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             attn_batches*q_seq_len/sequences);
+    }
+  }
+  assert(softmax_success);
+
+  if (is_training) {
+    apex_fused_dropout_cuda<at::Half,float,uint32_t>(
+                               static_cast<at::Half const*>(softmax_results.data_ptr()), 
+                               static_cast<at::Half*>(dropout_results.data_ptr()), 
+                               static_cast<uint8_t*>(dropout_mask.data_ptr()),
+                               dropout_elems,
+                               (1.0f - dropout_prob));
+  }
+ 
+  // Matmul2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr), 
+                             lead_dim_kv, 
+                             batch_stride_kv, 
+                             (is_training) ? static_cast<const half*>(dropout_results.data_ptr()) : static_cast<const half*>(softmax_results.data_ptr()) , 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             static_cast<half*>(matmul2_results.data_ptr()), 
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             attn_batches);
+
+  // Output Linear
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches_q, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(outputs.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO1_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return { 
+           input_lin_q_results, 
+           input_lin_kv_results, 
+           softmax_results,
+           dropout_results, 
+           dropout_mask, 
+           matmul2_results, 
+           outputs
+         };
+}
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_q_results,
+                               torch::Tensor const& input_lin_kv_results,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  ) 
+{
+  const int   embed_dim         = inputs_q.size(2);
+  const int   sequences         = inputs_q.size(1);
+  const int   q_seq_len         = inputs_q.size(0);
+  const int   k_seq_len         = inputs_kv.size(0);
+  const int   batches_q         = sequences * q_seq_len;
+  const int   batches_kv        = sequences * k_seq_len;
+  const int   head_dim          = embed_dim / heads;
+  const int   output_lin_q_dim  = embed_dim;
+  const int   output_lin_kv_dim = 2 * embed_dim;
+  const int   attn_batches      = heads * sequences;
+  const int   lead_dim_q        = attn_batches * head_dim;
+  const int   lead_dim_kv       = attn_batches * 2 *head_dim;
+  const int   batch_stride_q    = head_dim;
+  const int   batch_stride_kv   = 2 * head_dim;
+  const int   dropout_elems     = attn_batches * q_seq_len * k_seq_len;
+  const float alpha             = 1.0;
+  const float beta              = 0.0;
+  const float scale             = 1.0 / sqrt(static_cast<float>(head_dim));
+
+  // TODO: Streams can be used in Backprop but I haven't added more than one
+  // in my first attempt to create the code
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+ 
+  // Output Tensor Allocations
+  torch::Tensor input_q_grads          = torch::empty_like(inputs_q);
+  torch::Tensor input_kv_grads         = torch::empty_like(inputs_kv);
+  torch::Tensor input_weight_q_grads   = torch::empty_like(input_weights_q);
+  torch::Tensor input_weight_kv_grads  = torch::empty_like(input_weights_kv);
+  torch::Tensor output_weight_grads    = torch::empty_like(output_weights);
+  // Intermediate Tensor Allocations
+  at::Tensor output_lin_grads          = torch::empty_like(matmul2_results);
+  at::Tensor matmul2_grads             = torch::empty_like(dropout_results);
+  at::Tensor input_lin_q_output_grads  = torch::empty_like(input_lin_q_results);
+  at::Tensor input_lin_kv_output_grads = torch::empty_like(input_lin_kv_results);
+ 
+  auto q_lin_results_ptr = static_cast<half*>(input_lin_q_results.data_ptr());
+  auto k_lin_results_ptr = static_cast<half*>(input_lin_kv_results.data_ptr());
+  auto v_lin_results_ptr = static_cast<half*>(input_lin_kv_results.data_ptr()) + head_dim;
+  
+  auto q_lin_grads_ptr   = static_cast<half*>(input_lin_q_output_grads.data_ptr());
+  auto k_lin_grads_ptr   = static_cast<half*>(input_lin_kv_output_grads.data_ptr());
+  auto v_lin_grads_ptr   = static_cast<half*>(input_lin_kv_output_grads.data_ptr()) + head_dim;
+
+  char a_layout_n{'n'};
+  char a_layout_t{'t'};
+  char b_layout_n{'n'};
+  char b_layout_t{'t'}; 
+  
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+ 
+  // Output Linear Dgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches_q, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(output_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_lin_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+ 
+  // Output Linear Wgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             embed_dim,
+                             batches_q, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(output_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_weight_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // MatMul2 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr),
+                             lead_dim_kv, 
+                             batch_stride_kv,
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             beta, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len,
+                             attn_batches);
+  
+  // Matmul2 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             static_cast<const half*>(dropout_results.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             v_lin_grads_ptr, 
+                             lead_dim_kv, 
+                             batch_stride_kv, 
+                             attn_batches);
+
+  // Apply Dropout Mask and Scale by Dropout Probability 
+  apex_masked_scale_cuda<at::Half,float,uint32_t>(
+                             static_cast<at::Half const*>(matmul2_grads.data_ptr()),
+                             static_cast<at::Half*>(matmul2_grads.data_ptr()),
+                             static_cast<uint8_t const*>(dropout_mask.data_ptr()),
+                             dropout_elems,
+                             (1.0 / (1.0 - dropout_prob)));
+
+  // Softmax Grad
+  bool softmax_success = false;
+  softmax_success = dispatch_softmax_backward<half, half, float>(
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             reinterpret_cast<half const*>(softmax_results.data_ptr()),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  assert(softmax_success);
+
+  // Matmul1 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             scale, 
+                             k_lin_results_ptr, 
+                             lead_dim_kv, 
+                             batch_stride_kv, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             q_lin_grads_ptr, 
+                             lead_dim_q, 
+                             batch_stride_q, 
+                             attn_batches);
+  
+  // Matmul1 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             scale, 
+                             q_lin_results_ptr, 
+                             lead_dim_q, 
+                             batch_stride_q, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             k_lin_grads_ptr, 
+                             lead_dim_kv, 
+                             batch_stride_kv, 
+                             attn_batches);
+
+  // Input Linear Q Dgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim,
+                             batches_q, 
+                             output_lin_q_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights_q.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F, 
+                             output_lin_q_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_q_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO10_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // Input Linear Q Wgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             output_lin_q_dim,
+                             batches_q, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(inputs_q.data_ptr()),
+                             CUDA_R_16F,
+                             embed_dim,
+                             static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F,
+                             output_lin_q_dim,
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_weight_q_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // Input Linear KV Dgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim,
+                             batches_kv, 
+                             output_lin_kv_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights_kv.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(k_lin_grads_ptr),
+                             CUDA_R_16F, 
+                             output_lin_kv_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_kv_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO10_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // Input Linear KV Wgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             output_lin_kv_dim,
+                             batches_kv, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(inputs_kv.data_ptr()),
+                             CUDA_R_16F,
+                             embed_dim,
+                             static_cast<const void*>(k_lin_grads_ptr),
+                             CUDA_R_16F,
+                             output_lin_kv_dim,
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_weight_kv_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return { 
+           input_q_grads, 
+           input_kv_grads, 
+           input_weight_q_grads, 
+           input_weight_kv_grads, 
+           output_weight_grads
+         };
+}
+
+} // end namespace cublas_gemmex
+} // end namespace encdec 
+} // end namespace multihead_attn
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add.cpp
new file mode 100644
index 0000000000..314e9a9eee
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add.cpp
@@ -0,0 +1,198 @@
+#include <torch/extension.h>
+#include <vector>
+
+namespace multihead_attn {
+namespace encdec_norm_add {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,  
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+                               torch::Tensor const& lyr_nrm_gamma_weights,
+                               torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                                  );
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_q_results,
+                               torch::Tensor const& input_lin_kv_results,
+                               torch::Tensor const& lyr_nrm_results,
+                               torch::Tensor const& lyr_nrm_mean,
+                               torch::Tensor const& lyr_nrm_invvar,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+                               torch::Tensor const& lyr_nrm_gamma_weights,
+                               torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               torch::Tensor const& dropout_add_mask,
+                               float                dropout_prob
+                                                  );
+
+// C++ interface
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> fwd(
+                               bool                 use_mask,
+                               bool                 use_time_mask,
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+                               torch::Tensor const& lyr_nrm_gamma_weights,
+                               torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& pad_mask,
+                               float                dropout_prob
+                                                 )
+{
+  AT_ASSERTM(inputs_q.dim()               == 3, "expected 3D tensor");
+  AT_ASSERTM(inputs_kv.dim()              == 3, "expected 3D tensor");
+  AT_ASSERTM(lyr_nrm_gamma_weights.dim()  == 1, "expected 1D tensor");
+  AT_ASSERTM(lyr_nrm_beta_weights.dim()   == 1, "expected 1D tensor");
+  AT_ASSERTM(input_weights_q.dim()        == 2, "expected 2D tensor");
+  AT_ASSERTM(input_weights_kv.dim()       == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim()         == 2, "expected 2D tensor");
+
+  AT_ASSERTM(inputs_q.type().scalarType()              == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(inputs_kv.type().scalarType()             == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_gamma_weights.type().scalarType() == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_beta_weights.type().scalarType()  == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights_q.type().scalarType()       == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights_kv.type().scalarType()      == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType()        == at::ScalarType::Half, "Only HALF is supported");
+  
+  if (use_mask) {
+    AT_ASSERTM(pad_mask.dim()                     == 2,                    "expected 2D tensor");
+    AT_ASSERTM(pad_mask.type().scalarType()       == at::ScalarType::Byte, "Only BYTE is supported");
+  }
+  
+  return fwd_cuda(
+                                 use_time_mask,
+                                 is_training,
+                                 heads, 
+                                 inputs_q, 
+                                 inputs_kv,
+								 lyr_nrm_gamma_weights,
+								 lyr_nrm_beta_weights,
+                                 input_weights_q, 
+                                 input_weights_kv, 
+                                 output_weights, 
+                                 use_mask ? static_cast<const uint8_t*>(pad_mask.data_ptr()) : nullptr, 
+                                 dropout_prob
+                                );
+}
+
+std::vector<torch::Tensor> bwd(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_q_results,
+                               torch::Tensor const& input_lin_kv_results,
+                               torch::Tensor const& lyr_nrm_results,
+                               torch::Tensor const& lyr_nrm_mean,
+                               torch::Tensor const& lyr_nrm_invvar,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+							   torch::Tensor const& lyr_nrm_gamma_weights,
+							   torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               torch::Tensor const& dropout_add_mask,
+                               float                dropout_prob
+                                                  )
+{
+  AT_ASSERTM(output_grads.dim()          == 3, "expected 3D tensor");
+  AT_ASSERTM(matmul2_results.dim()       == 3, "expected 3D tensor");
+  AT_ASSERTM(dropout_results.dim()       == 3, "expected 3D tensor");
+  AT_ASSERTM(softmax_results.dim()       == 3, "expected 3D tensor");
+  AT_ASSERTM(input_lin_q_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(input_lin_kv_results.dim()  == 3, "expected 3D tensor");
+  AT_ASSERTM(lyr_nrm_results.dim()       == 3, "expected 3D tensor");
+  AT_ASSERTM(lyr_nrm_mean.dim()          == 1, "expected 1D tensor");
+  AT_ASSERTM(lyr_nrm_invvar.dim()        == 1, "expected 1D tensor");
+  AT_ASSERTM(inputs_q.dim()              == 3, "expected 3D tensor");
+  AT_ASSERTM(inputs_kv.dim()             == 3, "expected 3D tensor");
+  AT_ASSERTM(lyr_nrm_gamma_weights.dim() == 1, "expected 1D tensor");
+  AT_ASSERTM(lyr_nrm_beta_weights.dim()  == 1, "expected 1D tensor");
+  AT_ASSERTM(input_weights_q.dim()       == 2, "expected 2D tensor");
+  AT_ASSERTM(input_weights_kv.dim()      == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim()        == 2, "expected 2D tensor");
+  AT_ASSERTM(dropout_mask.dim()          == 3, "expected 3D tensor");
+  AT_ASSERTM(dropout_add_mask.dim()      == 3, "expected 3D tensor");
+  
+  AT_ASSERTM(output_grads.type().scalarType()          == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(matmul2_results.type().scalarType()       == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(dropout_results.type().scalarType()       == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(softmax_results.type().scalarType()       == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(input_lin_q_results.type().scalarType()   == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(input_lin_kv_results.type().scalarType()  == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_results.type().scalarType()       == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_mean.type().scalarType()          == at::ScalarType::Float, "Only FLOAT is supported");
+  AT_ASSERTM(lyr_nrm_invvar.type().scalarType()        == at::ScalarType::Float, "Only FLOAT is supported");
+  AT_ASSERTM(inputs_q.type().scalarType()              == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(inputs_kv.type().scalarType()             == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_gamma_weights.type().scalarType() == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_beta_weights.type().scalarType()  == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(input_weights_q.type().scalarType()       == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(input_weights_kv.type().scalarType()      == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType()        == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(dropout_mask.type().scalarType()          == at::ScalarType::Byte,  "Only BYTE is supported");
+  AT_ASSERTM(dropout_add_mask.type().scalarType()      == at::ScalarType::Byte,  "Only BYTE is supported");
+  
+  return bwd_cuda(
+                                 heads, 
+                                 output_grads,
+                                 matmul2_results,
+                                 dropout_results,
+                                 softmax_results, 
+                                 input_lin_q_results, 
+                                 input_lin_kv_results, 
+                                 lyr_nrm_results,
+                                 lyr_nrm_mean,
+                                 lyr_nrm_invvar,
+                                 inputs_q, 
+                                 inputs_kv, 
+								 lyr_nrm_gamma_weights,
+								 lyr_nrm_beta_weights,
+                                 input_weights_q,
+                                 input_weights_kv,
+                                 output_weights,
+                                 dropout_mask,
+                                 dropout_add_mask,
+                                 dropout_prob
+                                );
+}
+
+} // end namespace cublas_gemmex
+} // end namespace encdec_norm_add 
+} // end namespace multihead_attn
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &multihead_attn::encdec_norm_add::cublas_gemmex::fwd, "Encdec Multihead Attention Plus Layer Norm and Residual Add Forward.");
+  m.def("backward", &multihead_attn::encdec_norm_add::cublas_gemmex::bwd, "Encdec Multihead Attention Plus Layer Norm and Residual Add Backward.");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add_cuda.cu
new file mode 100644
index 0000000000..19d1ddd505
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add_cuda.cu
@@ -0,0 +1,640 @@
+#include <vector>
+#include <math.h>
+#include <iostream>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include "strided_batched_gemm.h"
+#include "softmax.h"
+#include "dropout.h"
+#include "layer_norm.h"
+
+// symbol to be automatically resolved by PyTorch libs
+extern THCState *state;
+
+namespace multihead_attn {
+namespace encdec_norm_add {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,
+							   bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+							   torch::Tensor const& lyr_nrm_gamma_weights,
+							   torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim         = inputs_q.size(2);
+  const int   sequences         = inputs_q.size(1);
+  const int   q_seq_len         = inputs_q.size(0);
+  const int   k_seq_len         = inputs_kv.size(0);
+  const int   batches_q         = sequences * q_seq_len;
+  const int   batches_kv        = sequences * k_seq_len;
+  const int   total_tokens_q    = batches_q * embed_dim;
+  const int   head_dim          = embed_dim / heads;
+  const int   output_lin_q_dim  = embed_dim;
+  const int   output_lin_kv_dim = 2 * embed_dim;
+  const int   attn_batches      = heads * sequences;
+  const int   lead_dim_q        = attn_batches * head_dim;
+  const int   lead_dim_kv       = attn_batches * 2 *head_dim;
+  const int   batch_stride_q    = head_dim;
+  const int   batch_stride_kv   = 2 * head_dim;
+  const int   dropout_elems     = attn_batches * q_seq_len * k_seq_len;
+  const float alpha             = 1.0;
+  const float beta              = 0.0;
+  const float scale             = 1.0 / sqrt(static_cast<float>(head_dim));
+ 
+  // There is no reason to use more than one stream as every kernel is 
+  // sequentially dependent
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // 3 Intermediate Results + Output (Note: dropout intermediates are generated by ATen library code)
+  auto act_options                   = inputs_q.options().requires_grad(false);
+  auto lyr_nrm_options               = act_options.dtype(torch::kFloat32);
+  auto mask_options                  = act_options.dtype(torch::kUInt8);
+  
+  torch::Tensor lyr_nrm_mean         = torch::empty({batches_q}, lyr_nrm_options);
+  torch::Tensor lyr_nrm_invvar       = torch::empty({batches_q}, lyr_nrm_options);
+  torch::Tensor lyr_nrm_results      = torch::empty_like(inputs_q, act_options);
+
+  torch::Tensor input_lin_q_results  = torch::empty({q_seq_len, sequences, output_lin_q_dim},  act_options);
+  torch::Tensor input_lin_kv_results = torch::empty({k_seq_len, sequences, output_lin_kv_dim}, act_options);
+  torch::Tensor softmax_results      = torch::empty({attn_batches, q_seq_len, k_seq_len},      act_options);
+  torch::Tensor dropout_results      = torch::empty({attn_batches, q_seq_len, k_seq_len},      act_options);
+  torch::Tensor dropout_mask         = torch::empty({attn_batches, q_seq_len, k_seq_len},      mask_options);
+  torch::Tensor matmul2_results      = torch::empty({q_seq_len, attn_batches, head_dim},       act_options);
+  torch::Tensor output_lin_results   = torch::empty_like(inputs_q, act_options);
+  torch::Tensor dropout_add_mask     = torch::empty_like(inputs_q, mask_options);
+  torch::Tensor outputs              = torch::empty_like(inputs_q, act_options);
+
+  // Input Linear Results Pointers to Q, K, and V of interviewed activations
+  void* q_lin_results_ptr   = static_cast<void*>(input_lin_q_results.data_ptr());
+  void* k_lin_results_ptr   = static_cast<void*>(input_lin_kv_results.data_ptr());
+  void* v_lin_results_ptr   = static_cast<void*>(static_cast<half*>(input_lin_kv_results.data_ptr()) + head_dim);
+
+  // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)
+  void* softmax_results_ptr = static_cast<void*>(softmax_results.data_ptr());
+  
+  char a_layout_t{'t'};
+  char a_layout_n{'n'};
+  char b_layout_n{'n'};
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+  // Layer Norm
+  HostApplyLayerNorm<at::Half,float>(
+                             static_cast<at::Half*>(lyr_nrm_results.data_ptr()),
+                             static_cast<float*>(lyr_nrm_mean.data_ptr()),
+                             static_cast<float*>(lyr_nrm_invvar.data_ptr()),
+                             static_cast<const at::Half*>(inputs_q.data_ptr()),
+                             static_cast<int>(batches_q), // n1
+                             static_cast<int>(embed_dim), // n2
+                             1.0e-5,
+							 static_cast<const at::Half*>(lyr_nrm_gamma_weights.data_ptr()),
+							 static_cast<const at::Half*>(lyr_nrm_beta_weights.data_ptr()));
+
+  // Input Linear Q Fwd
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             output_lin_q_dim, 
+                             batches_q, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights_q.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             //static_cast<const void*>(inputs_q.data_ptr()),
+                             static_cast<const void*>(lyr_nrm_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             q_lin_results_ptr,
+                             CUDA_R_16F, 
+                             output_lin_q_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // Input Linear KV Fwd
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             output_lin_kv_dim, 
+                             batches_kv, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights_kv.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(inputs_kv.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             k_lin_results_ptr,
+                             CUDA_R_16F, 
+                             output_lin_kv_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             scale, 
+                             static_cast<const half*>(k_lin_results_ptr), 
+                             lead_dim_kv, 
+                             batch_stride_kv, 
+                             static_cast<const half*>(q_lin_results_ptr),
+                             lead_dim_q, 
+                             batch_stride_q, 
+                             beta, 
+                             static_cast<half*>(softmax_results_ptr), 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             attn_batches);
+
+  // Padded Softmax
+  bool softmax_success = false;
+  if (pad_mask == nullptr) {
+    softmax_success = dispatch_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  } else {
+    if (use_time_mask) {
+      softmax_success = dispatch_time_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             q_seq_len);
+    } else {
+      softmax_success = dispatch_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             attn_batches*q_seq_len/sequences);
+    }
+  }
+  assert(softmax_success);
+  
+  if (is_training) {
+    apex_fused_dropout_cuda<at::Half,float,uint32_t>(
+                             static_cast<at::Half const*>(softmax_results.data_ptr()), 
+                             static_cast<at::Half*>(dropout_results.data_ptr()), 
+                             static_cast<uint8_t*>(dropout_mask.data_ptr()),
+                             dropout_elems,
+                             (1.0f - dropout_prob));
+  }
+
+  // Matmul2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr), 
+                             lead_dim_kv, 
+                             batch_stride_kv, 
+                             (is_training) ? static_cast<const half*>(dropout_results.data_ptr()) : static_cast<const half*>(softmax_results.data_ptr()), 
+                             //static_cast<const half*>(dropout_results.data_ptr()), 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             static_cast<half*>(matmul2_results.data_ptr()), 
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             attn_batches);
+
+  // Output Linear
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches_q, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_lin_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO1_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+ 
+  // End-of-block Dropout-Add 
+  if (is_training) {
+    apex_dropout_add_cuda<at::Half,float,uint32_t>(
+                             static_cast<at::Half const*>(output_lin_results.data_ptr()), 
+                             static_cast<at::Half const*>(inputs_q.data_ptr()), 
+                             static_cast<at::Half*>(outputs.data_ptr()), 
+                             static_cast<uint8_t*>(dropout_add_mask.data_ptr()),
+                             total_tokens_q,
+                             (1.0f - dropout_prob));
+  } else {
+    apex_add_cuda<at::Half,float,uint32_t>(
+                             static_cast<at::Half const*>(output_lin_results.data_ptr()), 
+                             static_cast<at::Half const*>(inputs_q.data_ptr()), 
+                             static_cast<at::Half*>(outputs.data_ptr()), 
+                             total_tokens_q);
+  }
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return {
+		   lyr_nrm_results,
+		   lyr_nrm_mean,
+		   lyr_nrm_invvar, 
+           input_lin_q_results, 
+           input_lin_kv_results, 
+           softmax_results,
+           dropout_results, 
+           dropout_mask, 
+           matmul2_results, 
+		   dropout_add_mask, 
+           outputs
+         };
+}
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_q_results,
+                               torch::Tensor const& input_lin_kv_results,
+                               torch::Tensor const& lyr_nrm_results,
+                               torch::Tensor const& lyr_nrm_mean,
+                               torch::Tensor const& lyr_nrm_invvar,
+                               torch::Tensor const& inputs_q, 
+                               torch::Tensor const& inputs_kv, 
+							   torch::Tensor const& lyr_nrm_gamma_weights,
+							   torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights_q,
+                               torch::Tensor const& input_weights_kv,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               torch::Tensor const& dropout_add_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim         = inputs_q.size(2);
+  const int   sequences         = inputs_q.size(1);
+  const int   q_seq_len         = inputs_q.size(0);
+  const int   k_seq_len         = inputs_kv.size(0);
+  const int   batches_q         = sequences * q_seq_len;
+  const int   batches_kv        = sequences * k_seq_len;
+  const int   total_tokens_q    = batches_q * embed_dim;
+  const int   head_dim          = embed_dim / heads;
+  const int   output_lin_q_dim  = embed_dim;
+  const int   output_lin_kv_dim = 2 * embed_dim;
+  const int   attn_batches      = heads * sequences;
+  const int   lead_dim_q        = attn_batches * head_dim;
+  const int   lead_dim_kv       = attn_batches * 2 *head_dim;
+  const int   batch_stride_q    = head_dim;
+  const int   batch_stride_kv   = 2 * head_dim;
+  const int   dropout_elems     = attn_batches * q_seq_len * k_seq_len;
+  const float alpha             = 1.0;
+  const float beta              = 0.0;
+  const float scale             = 1.0 / sqrt(static_cast<float>(head_dim));
+
+  // TODO: Streams can be used in Backprop but I haven't added more than one
+  // in my first attempt to create the code
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+ 
+  // Output Tensor Allocations
+  torch::Tensor input_q_grads          = torch::empty_like(inputs_q);
+  torch::Tensor input_kv_grads         = torch::empty_like(inputs_kv);
+  torch::Tensor lyr_nrm_gamma_grads    = torch::empty_like(lyr_nrm_gamma_weights);
+  torch::Tensor lyr_nrm_beta_grads     = torch::empty_like(lyr_nrm_beta_weights);
+  torch::Tensor input_weight_q_grads   = torch::empty_like(input_weights_q);
+  torch::Tensor input_weight_kv_grads  = torch::empty_like(input_weights_kv);
+  torch::Tensor output_weight_grads    = torch::empty_like(output_weights);
+  // Intermediate Tensor Allocations
+  at::Tensor dropout_add_grads         = torch::empty_like(output_grads);
+  at::Tensor output_lin_grads          = torch::empty_like(matmul2_results);
+  at::Tensor matmul2_grads             = torch::empty_like(dropout_results);
+  at::Tensor input_lin_q_output_grads  = torch::empty_like(input_lin_q_results);
+  at::Tensor input_lin_kv_output_grads = torch::empty_like(input_lin_kv_results);
+  at::Tensor input_lin_q_grads         = torch::empty_like(inputs_q);
+ 
+  auto q_lin_results_ptr = static_cast<half*>(input_lin_q_results.data_ptr());
+  auto k_lin_results_ptr = static_cast<half*>(input_lin_kv_results.data_ptr());
+  auto v_lin_results_ptr = static_cast<half*>(input_lin_kv_results.data_ptr()) + head_dim;
+  
+  auto q_lin_grads_ptr   = static_cast<half*>(input_lin_q_output_grads.data_ptr());
+  auto k_lin_grads_ptr   = static_cast<half*>(input_lin_kv_output_grads.data_ptr());
+  auto v_lin_grads_ptr   = static_cast<half*>(input_lin_kv_output_grads.data_ptr()) + head_dim;
+
+  char a_layout_n{'n'};
+  char a_layout_t{'t'};
+  char b_layout_n{'n'};
+  char b_layout_t{'t'}; 
+  
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+  
+  // Dropout Add Backward  
+  apex_masked_scale_cuda<at::Half,float,uint32_t>(
+                             static_cast<at::Half const*>(output_grads.data_ptr()),
+							 static_cast<at::Half*>(dropout_add_grads.data_ptr()),
+							 static_cast<uint8_t const*>(dropout_add_mask.data_ptr()),
+							 total_tokens_q,
+                             (1.0 / (1.0 - dropout_prob)));
+ 
+  // Output Linear Dgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches_q, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(dropout_add_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_lin_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+ 
+  // Output Linear Wgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             embed_dim,
+                             batches_q, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(dropout_add_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_weight_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // MatMul2 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr),
+                             lead_dim_kv, 
+                             batch_stride_kv,
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             beta, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len,
+                             attn_batches);
+  
+  // Matmul2 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             static_cast<const half*>(dropout_results.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             v_lin_grads_ptr, 
+                             lead_dim_kv, 
+                             batch_stride_kv, 
+                             attn_batches);
+
+  // Apply Dropout Mask and Scale by Dropout Probability 
+  apex_masked_scale_cuda<at::Half,float,uint32_t>(
+                             static_cast<at::Half const*>(matmul2_grads.data_ptr()),
+							 static_cast<at::Half*>(matmul2_grads.data_ptr()),
+							 static_cast<uint8_t const*>(dropout_mask.data_ptr()),
+							 dropout_elems,
+                             (1.0 / (1.0 - dropout_prob)));
+
+  // Softmax Grad
+  bool softmax_success = false;
+  softmax_success = dispatch_softmax_backward<half, half, float>(
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             reinterpret_cast<half const*>(softmax_results.data_ptr()),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  assert(softmax_success);
+
+  // Matmul1 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             scale, 
+                             k_lin_results_ptr, 
+                             lead_dim_kv, 
+                             batch_stride_kv, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             q_lin_grads_ptr, 
+                             lead_dim_q, 
+                             batch_stride_q, 
+                             attn_batches);
+  
+  // Matmul1 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             scale, 
+                             q_lin_results_ptr, 
+                             lead_dim_q, 
+                             batch_stride_q, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             k_lin_grads_ptr, 
+                             lead_dim_kv, 
+                             batch_stride_kv, 
+                             attn_batches);
+
+  // Input Linear Q Dgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim,
+                             batches_q, 
+                             output_lin_q_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights_q.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F, 
+                             output_lin_q_dim, 
+                             static_cast<const void*>(&beta),
+                             //static_cast<void*>(input_q_grads.data_ptr()),
+                             static_cast<void*>(input_lin_q_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO10_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // Input Linear Q Wgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             output_lin_q_dim,
+                             batches_q, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(inputs_q.data_ptr()),
+                             CUDA_R_16F,
+                             embed_dim,
+                             static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F,
+                             output_lin_q_dim,
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_weight_q_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // Input Linear KV Dgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim,
+                             batches_kv, 
+                             output_lin_kv_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights_kv.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(k_lin_grads_ptr),
+                             CUDA_R_16F, 
+                             output_lin_kv_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_kv_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO10_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // Input Linear KV Wgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             output_lin_kv_dim,
+                             batches_kv, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(inputs_kv.data_ptr()),
+                             CUDA_R_16F,
+                             embed_dim,
+                             static_cast<const void*>(k_lin_grads_ptr),
+                             CUDA_R_16F,
+                             output_lin_kv_dim,
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_weight_kv_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+ 
+  // Fused Layer Norm Bwd with Residual Add
+  HostLayerNormGradient<half,float>(
+                             static_cast<const half*>(input_lin_q_grads.data_ptr()),
+                             static_cast<half const*>(output_grads.data_ptr()), 
+                             static_cast<const float*>(lyr_nrm_mean.data_ptr()),
+                             static_cast<const float*>(lyr_nrm_invvar.data_ptr()),
+                             inputs_q,
+                             static_cast<int>(batches_q),  // n1
+                             static_cast<int>(embed_dim),  // n2
+                             static_cast<const half*>(lyr_nrm_gamma_weights.data_ptr()),
+                             static_cast<const half*>(lyr_nrm_beta_weights.data_ptr()),
+                             1.0e-5,
+                             static_cast<half*>(input_q_grads.data_ptr()),
+                             static_cast<half*>(lyr_nrm_gamma_grads.data_ptr()),
+                             static_cast<half*>(lyr_nrm_beta_grads.data_ptr())
+                                  );
+  
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return { 
+           input_q_grads, 
+           input_kv_grads, 
+           lyr_nrm_gamma_grads, 
+           lyr_nrm_beta_grads, 
+           input_weight_q_grads, 
+           input_weight_kv_grads, 
+           output_weight_grads
+         };
+}
+
+} // end namespace cublas_gemmex
+} // end namespace encdec_norm_add 
+} // end namespace multihead_attn
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/layer_norm.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/layer_norm.h
new file mode 100644
index 0000000000..07c8aac970
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/layer_norm.h
@@ -0,0 +1,740 @@
+#include "ATen/ATen.h"
+#include "ATen/cuda/DeviceUtils.cuh"
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+template<typename U> __device__
+void cuWelfordOnlineSum(
+  const U curr,
+  U& mu,
+  U& sigma2,
+  U& count)
+{
+  count = count + U(1);
+  U delta = curr - mu;
+  U lmean = mu + delta / count;
+  mu = lmean;
+  U delta2 = curr - lmean;
+  sigma2 = sigma2 + delta * delta2;
+}
+
+template<typename U> __device__
+void cuChanOnlineSum(
+  const U muB,
+  const U sigma2B,
+  const U countB,
+  U& mu,
+  U& sigma2,
+  U& count)
+{
+  U delta = muB - mu;
+  U nA = count;
+  U nB = countB;
+  count = count + countB;
+  U nX = count;
+  if (nX > U(0)) {
+    nA = nA / nX;
+    nB = nB / nX;
+    mu = nA*mu + nB*muB;
+    sigma2 = sigma2 + sigma2B + delta * delta * nA * nB * nX;
+  } else {
+    mu = U(0);
+    sigma2 = U(0);
+  }
+}
+
+template<typename T, typename U> __device__
+void cuWelfordMuSigma2(
+  const T* __restrict__ vals,
+  const int n1,
+  const int n2,
+  const int i1,
+  U& mu,
+  U& sigma2,
+  U* buf) 
+{
+  // Assumptions:
+  // 1) blockDim.x == warpSize
+  // 2) Tensor is contiguous
+  // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.
+  //
+  // compute variance and mean over n2
+  U count = U(0);
+  mu= U(0);
+  sigma2 = U(0);
+  if (i1 < n1) {
+    // one warp normalizes one n1 index,
+    // synchronization is implicit
+    // initialize with standard Welford algorithm
+    const int numx = blockDim.x * blockDim.y;
+    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
+    const T* lvals = vals + i1*n2;
+    int l = 4*thrx;
+    for (;  l+3 < n2;  l+=4*numx) {
+      for (int k = 0;  k < 4;  ++k) {
+        U curr = static_cast<U>(lvals[l+k]);
+        cuWelfordOnlineSum<U>(curr,mu,sigma2,count);
+      }
+    }
+    for (;  l < n2;  ++l) {
+      U curr = static_cast<U>(lvals[l]);
+      cuWelfordOnlineSum<U>(curr,mu,sigma2,count);
+    }
+    // intra-warp reductions
+    for (int l = 0;  l <= 4;  ++l) {
+      int srcLaneB = (threadIdx.x+(1<<l))&31;
+      U muB = WARP_SHFL(mu, srcLaneB);
+      U countB = WARP_SHFL(count, srcLaneB);
+      U sigma2B = WARP_SHFL(sigma2, srcLaneB);
+      cuChanOnlineSum<U>(muB,sigma2B,countB,mu,sigma2,count);
+    }
+    // threadIdx.x == 0 has correct values for each warp
+    // inter-warp reductions
+    if (blockDim.y > 1) {
+      U* ubuf = (U*)buf;
+      U* ibuf = (U*)(ubuf + blockDim.y);
+      for (int offset = blockDim.y/2;  offset > 0;  offset /= 2) {
+        // upper half of warps write to shared
+        if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2*offset) {
+          const int wrt_y = threadIdx.y - offset;
+          ubuf[2*wrt_y] = mu;
+          ubuf[2*wrt_y+1] = sigma2;
+          ibuf[wrt_y] = count;
+        }
+        __syncthreads();
+        // lower half merges
+        if (threadIdx.x == 0 && threadIdx.y < offset) {
+          U muB = ubuf[2*threadIdx.y];
+          U sigma2B = ubuf[2*threadIdx.y+1];
+          U countB = ibuf[threadIdx.y];
+          cuChanOnlineSum<U>(muB,sigma2B,countB,mu,sigma2,count);
+        }
+        __syncthreads();
+      }
+      // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values
+      if (threadIdx.x == 0 && threadIdx.y == 0) {
+        ubuf[0] = mu;
+        ubuf[1] = sigma2;
+      }
+      __syncthreads();
+      mu = ubuf[0];
+      sigma2 = ubuf[1]/U(n2);
+      // don't care about final value of count, we know count == n2
+    } else {
+      mu = WARP_SHFL(mu, 0);
+      sigma2 = WARP_SHFL(sigma2/U(n2), 0);
+    }
+  }
+}
+
+template<> __device__
+void cuWelfordMuSigma2(
+  const at::Half* __restrict__ vals,
+  const int n1,
+  const int n2,
+  const int i1,
+  float& mu,
+  float& sigma2,
+  float* buf) 
+{
+  // Assumptions:
+  // 1) blockDim.x == warpSize
+  // 2) Tensor is contiguous
+  // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.
+  //
+  // compute variance and mean over n2
+  float count = 0.0f;
+  mu= float(0);
+  sigma2 = float(0);
+
+  if (i1 < n1) {
+    // one warp normalizes one n1 index,
+    // synchronization is implicit
+    // initialize with standard Welford algorithm
+    const int numx = blockDim.x * blockDim.y;
+    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
+    const at::Half* lvals = vals + i1*n2;
+    int l = 8*thrx;
+    if ((((size_t)lvals)&3) != 0) {
+      // 16 bit alignment
+      // first thread consumes first point
+      if (thrx == 0) {
+        float curr = static_cast<float>(lvals[0]);
+        cuWelfordOnlineSum(curr,mu,sigma2,count);
+      }
+      ++l;
+    }
+    // at this point, lvals[l] are 32 bit aligned for all threads.
+    for (;  l+7 < n2;  l+=8*numx) {
+      for (int k = 0;  k < 8;  k+=2) {
+        float2 curr = __half22float2(*((__half2*)(lvals+l+k)));
+        cuWelfordOnlineSum(curr.x,mu,sigma2,count);
+	cuWelfordOnlineSum(curr.y,mu,sigma2,count);
+      }
+    }
+    for (;  l < n2;  ++l) {
+      float curr = static_cast<float>(lvals[l]);
+      cuWelfordOnlineSum(curr,mu,sigma2,count);
+    }
+    // intra-warp reductions
+    for (int l = 0;  l <= 4;  ++l) {
+      int srcLaneB = (threadIdx.x+(1<<l))&31;
+      float muB = WARP_SHFL(mu, srcLaneB);
+      float countB = WARP_SHFL(count, srcLaneB);
+      float sigma2B = WARP_SHFL(sigma2, srcLaneB);
+      cuChanOnlineSum(muB,sigma2B,countB,mu,sigma2,count);
+    }
+    // threadIdx.x == 0 has correct values for each warp
+    // inter-warp reductions
+    if (blockDim.y > 1) {
+      float* ubuf = (float*)buf;
+      float* ibuf = (float*)(ubuf + blockDim.y);
+      for (int offset = blockDim.y/2;  offset > 0;  offset /= 2) {
+        // upper half of warps write to shared
+        if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2*offset) {
+          const int wrt_y = threadIdx.y - offset;
+          ubuf[2*wrt_y] = mu;
+          ubuf[2*wrt_y+1] = sigma2;
+          ibuf[wrt_y] = count;
+        }
+        __syncthreads();
+        // lower half merges
+        if (threadIdx.x == 0 && threadIdx.y < offset) {
+          float muB = ubuf[2*threadIdx.y];
+          float sigma2B = ubuf[2*threadIdx.y+1];
+          float countB = ibuf[threadIdx.y];
+          cuChanOnlineSum(muB,sigma2B,countB,mu,sigma2,count);
+        }
+        __syncthreads();
+      }
+      // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values
+      if (threadIdx.x == 0 && threadIdx.y == 0) {
+        ubuf[0] = mu;
+        ubuf[1] = sigma2;
+      }
+      __syncthreads();
+      mu = ubuf[0];
+      sigma2 = ubuf[1]/float(n2);
+      // don't care about final value of count, we know count == n2
+    } else {
+      mu = WARP_SHFL(mu, 0);
+      sigma2 = WARP_SHFL(sigma2/float(n2), 0);
+    }
+  }
+}
+
+template<typename U> U rsqrt(U v) {
+  return U(1) / sqrt(v);
+}
+template<> float rsqrt(float v) {
+  return rsqrtf(v);
+}
+template<> double rsqrt(double v) {
+  return rsqrt(v);
+}
+
+namespace {
+// This is the un-specialized struct.  Note that we prevent instantiation of this
+// struct by putting an undefined symbol in the function body so it won't compile.
+//  template <typename T>
+//  struct SharedMemory
+//  {
+//      // Ensure that we won't compile any un-specialized types
+//      __device__ T *getPointer()
+//      {
+//          extern __device__ void error(void);
+//          error();
+//          return NULL;
+//      }
+//  };
+// https://github.com/NVIDIA/apex/issues/246
+template <typename T>
+struct SharedMemory;
+
+template <>
+struct SharedMemory <float>
+{
+    __device__ float *getPointer()
+    {
+        extern __shared__ float s_float[];
+        return s_float;
+    }
+};
+
+template <>
+struct SharedMemory <double>
+{
+    __device__ double *getPointer()
+    {
+        extern __shared__ double s_double[];
+        return s_double;
+    }
+};
+}
+
+template<typename T, typename U> __global__
+void cuApplyLayerNorm(
+  T* __restrict__ output_vals,
+  U* __restrict__ mean,
+  U* __restrict__ invvar,
+  const T* __restrict__ vals,
+  const int n1,
+  const int n2,
+  const U epsilon,
+  const T* __restrict__ gamma,
+  const T* __restrict__ beta
+  ) 
+{
+  // Assumptions:
+  // 1) blockDim.x == warpSize
+  // 2) Tensors are contiguous
+  //
+  for (auto i1=blockIdx.y; i1 < n1; i1 += gridDim.y) {
+    SharedMemory<U> shared;
+    U* buf = shared.getPointer();
+    U mu,sigma2;
+    cuWelfordMuSigma2(vals,n1,n2,i1,mu,sigma2,buf);
+    const T* lvals = vals + i1*n2;
+    T* ovals = output_vals + i1*n2;
+    U c_invvar = rsqrt(sigma2 + epsilon);
+    const int numx = blockDim.x * blockDim.y;
+    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
+    if (gamma != NULL && beta != NULL) {
+      for (int i = thrx;  i < n2;  i+=numx) {
+        U curr = static_cast<U>(lvals[i]);
+        ovals[i] = gamma[i] * static_cast<T>(c_invvar * (curr - mu)) + beta[i];
+      }
+    } else {
+      for (int i = thrx;  i < n2;  i+=numx) {
+        U curr = static_cast<U>(lvals[i]);
+        ovals[i] = static_cast<T>(c_invvar * (curr - mu));
+      }
+    }
+    if (threadIdx.x == 0 && threadIdx.y == 0) {
+      mean[i1] = mu;
+      invvar[i1] = c_invvar;
+    }
+  }
+}
+
+template<typename T, typename U> __device__
+void cuLoadWriteStridedInputs(
+    const int i1_block,
+    const int thr_load_row_off,
+    const int thr_load_col_off,
+    const int i2_off,
+    const int row_stride,
+    U* warp_buf1,
+    U* warp_buf2,
+    const T* input,
+    const T* dout,
+    const int i1_end,
+    const int n2,
+    const U* __restrict__ mean,
+    const U* __restrict__ invvar
+    )
+{
+  int i1 = i1_block+thr_load_row_off;
+  if (i1 < i1_end) {
+    U curr_mean = mean[i1];
+    U curr_invvar = invvar[i1];
+    for (int k = 0;  k < blockDim.y;  ++k) {
+      int i2 = i2_off + k;
+      int load_idx = i1*n2+i2;
+      int write_idx = thr_load_row_off*row_stride+thr_load_col_off+k;
+      if (i2<n2) {
+        U curr_input = static_cast<U>(input[load_idx]);
+	U curr_dout = static_cast<U>(dout[load_idx]);
+	warp_buf1[write_idx] = curr_dout;
+	warp_buf2[write_idx] = curr_dout * (curr_input - curr_mean) * curr_invvar;
+      } else {
+        warp_buf1[write_idx] = U(0);
+        warp_buf2[write_idx] = U(0);
+      }
+    }
+  } else {
+    for (int k = 0;  k < blockDim.y;  ++k) {
+      int write_idx = thr_load_row_off*row_stride+thr_load_col_off+k;
+      warp_buf1[write_idx] = U(0);
+      warp_buf2[write_idx] = U(0);
+    }
+  }
+}
+
+template<typename T, typename U> __device__
+void cuLoadAddStridedInputs(
+    const int i1_block,
+    const int thr_load_row_off,
+    const int thr_load_col_off,
+    const int i2_off,
+    const int row_stride,
+    U* warp_buf1,
+    U* warp_buf2,
+    const T* input,
+    const T* dout,
+    const int i1_end,
+    const int n2,
+    const U* __restrict__ mean,
+    const U* __restrict__ invvar
+    )
+{
+  int i1 = i1_block+thr_load_row_off;
+  if (i1 < i1_end) {
+    U curr_mean = mean[i1];
+    U curr_invvar = invvar[i1];
+    for (int k = 0;  k < blockDim.y;  ++k) {
+      int i2 = i2_off + k;
+      int load_idx = i1*n2+i2;
+      int write_idx = thr_load_row_off*row_stride+thr_load_col_off+k;
+      if (i2<n2) {
+        U curr_input = static_cast<U>(input[load_idx]);
+	U curr_dout = static_cast<U>(dout[load_idx]);
+	warp_buf1[write_idx] += curr_dout;
+	warp_buf2[write_idx] += curr_dout * (curr_input - curr_mean) * curr_invvar;
+      }
+    }
+  }
+}
+
+template<typename T, typename U> __global__
+void cuComputePartGradGammaBeta(
+    const T* __restrict__ dout,
+    const T* __restrict__ input,
+    const int n1,
+    const int n2,
+    const U* __restrict__ mean,
+    const U* __restrict__ invvar,
+    U epsilon,
+    U* part_grad_gamma,
+    U* part_grad_beta)
+{
+    const int numsegs_n1 = (n1+blockDim.y*blockDim.y-1) / (blockDim.y*blockDim.y);
+    const int segs_per_block = (numsegs_n1 + gridDim.y - 1) / gridDim.y;
+    const int i1_beg = blockIdx.y * segs_per_block * blockDim.y*blockDim.y;
+    const int i1_beg_plus_one = (blockIdx.y+1) * segs_per_block * blockDim.y*blockDim.y;
+    const int i1_end = i1_beg_plus_one < n1 ? i1_beg_plus_one : n1;
+    const int row_stride = blockDim.x+1;
+    const int thr_load_col_off = (threadIdx.x*blockDim.y)&(blockDim.x-1);
+    const int thr_load_row_off = (threadIdx.x*blockDim.y)/blockDim.x + threadIdx.y*blockDim.y;
+    const int i2_off = blockIdx.x * blockDim.x + thr_load_col_off;
+    SharedMemory<U> shared;
+    U* buf = shared.getPointer(); // buf has at least blockDim.x * blockDim.y * blockDim.y + (blockDim.y - 1)*(blockDim.x/blockDim.y) elements
+    U* warp_buf1 = (U*)buf;
+    U* warp_buf2 = warp_buf1 + blockDim.y * blockDim.y * row_stride;
+    // compute partial sums from strided inputs
+    // do this to increase number of loads in flight
+    cuLoadWriteStridedInputs(i1_beg,thr_load_row_off,thr_load_col_off,i2_off,row_stride,warp_buf1,warp_buf2,input,dout,i1_end,n2,mean,invvar);
+    for (int i1_block = i1_beg+blockDim.y*blockDim.y;  i1_block < i1_end;  i1_block+=blockDim.y*blockDim.y) {
+      cuLoadAddStridedInputs(i1_block,thr_load_row_off,thr_load_col_off,i2_off,row_stride,warp_buf1,warp_buf2,input,dout,i1_end,n2,mean,invvar);
+    }
+    __syncthreads();
+    // inter-warp reductions
+    // sum within each warp
+    U acc1 = U(0);
+    U acc2 = U(0);
+    for (int k = 0;  k < blockDim.y;  ++k) {
+      int row1 = threadIdx.y + k*blockDim.y;
+      int idx1 = row1*row_stride + threadIdx.x;
+      acc1 += warp_buf1[idx1];
+      acc2 += warp_buf2[idx1];
+    }
+    warp_buf1[threadIdx.y*row_stride+threadIdx.x] = acc1;
+    warp_buf2[threadIdx.y*row_stride+threadIdx.x] = acc2;
+    __syncthreads();
+    // sum all warps
+    for (int offset = blockDim.y/2;  offset > 1;  offset /= 2) {
+      if (threadIdx.y < offset) {
+        int row1 = threadIdx.y;
+	int row2 = threadIdx.y + offset;
+	int idx1 = row1*row_stride + threadIdx.x;
+	int idx2 = row2*row_stride + threadIdx.x;
+	warp_buf1[idx1] += warp_buf1[idx2];
+	warp_buf2[idx1] += warp_buf2[idx2];
+      }
+      __syncthreads();
+    }
+    int i2 = blockIdx.x * blockDim.x + threadIdx.x;
+    if (threadIdx.y == 0 && i2 < n2) {
+      int row1 = threadIdx.y;
+      int row2 = threadIdx.y + 1;
+      int idx1 = row1*row_stride + threadIdx.x;
+      int idx2 = row2*row_stride + threadIdx.x;
+      part_grad_beta[blockIdx.y*n2+i2] = warp_buf1[idx1] + warp_buf1[idx2];
+      part_grad_gamma[blockIdx.y*n2+i2] = warp_buf2[idx1] + warp_buf2[idx2];
+    }
+}
+
+template<typename T, typename U> __global__
+void cuComputeGradGammaBeta(
+    const U* part_grad_gamma,
+    const U* part_grad_beta,
+    const int part_size,
+    const int n1,
+    const int n2,
+    T* grad_gamma,
+    T* grad_beta)
+{
+    // sum partial gradients for gamma and beta
+    SharedMemory<U> shared;
+    U* buf = shared.getPointer(); 
+    int i2 = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i2 < n2) {
+      // each warp does sequential reductions until reduced part_size is num_warps
+      int num_warp_reductions = part_size / blockDim.y;
+      U sum_gamma = U(0);
+      U sum_beta = U(0);
+      const U* part_grad_gamma_ptr = part_grad_gamma + threadIdx.y * num_warp_reductions * n2 + i2;
+      const U* part_grad_beta_ptr = part_grad_beta + threadIdx.y * num_warp_reductions * n2 + i2;
+      for (int warp_offset = 0;  warp_offset < num_warp_reductions;  ++warp_offset) {
+        sum_gamma += part_grad_gamma_ptr[warp_offset*n2];
+        sum_beta += part_grad_beta_ptr[warp_offset*n2];
+      }
+      // inter-warp reductions
+      const int nbsize3 = blockDim.x * blockDim.y / 2;
+      for (int offset = blockDim.y/2;  offset >= 1;  offset /= 2) {
+        // top half write to shared memory
+        if (threadIdx.y >= offset && threadIdx.y < 2*offset) {
+          const int write_idx = (threadIdx.y - offset) * blockDim.x + threadIdx.x;
+          buf[write_idx] = sum_gamma;
+          buf[write_idx+nbsize3] = sum_beta;
+        }
+        __syncthreads();
+        // bottom half sums
+        if (threadIdx.y < offset) {
+          const int read_idx = threadIdx.y * blockDim.x + threadIdx.x;
+          sum_gamma += buf[read_idx];
+          sum_beta += buf[read_idx+nbsize3];
+        }
+        __syncthreads();
+      }
+      // write out fully summed gradients
+      if (threadIdx.y == 0) {
+        grad_gamma[i2] = sum_gamma;
+        grad_beta[i2] = sum_beta;
+      }
+    }
+}
+
+template<typename T, typename U> __global__
+void cuComputeGradInput(
+    const T* __restrict__ dout,
+    const T* __restrict__ dout_resid,
+    const T* __restrict__ input,
+    const int n1,
+    const int n2,
+    const U* __restrict__ mean,
+    const U* __restrict__ invvar,
+    U epsilon,
+    const T* gamma,
+    T* grad_input)
+{
+  for (auto i1=blockIdx.y; i1 < n1; i1 += gridDim.y) {
+    U sum_loss1 = U(0);
+    U sum_loss2 = U(0);
+    const U c_mean = mean[i1];
+    const U c_invvar = invvar[i1];
+    const T* k_input = input + i1*n2;
+    const T* k_dout = dout + i1*n2;
+    const T* k_dout_resid = dout_resid + i1*n2;
+    const int numx = blockDim.x * blockDim.y;
+    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
+    if (gamma != NULL) {
+      int l = 4*thrx;
+      for (;  l+3 < n2;  l+=4*numx) {
+        for (int k = 0;  k < 4;  ++k) {
+          const U c_h = static_cast<U>(k_input[l+k]);
+          const U c_loss = static_cast<U>(k_dout[l+k]);
+          sum_loss1 += c_loss * static_cast<U>(gamma[l+k]);
+          sum_loss2 += c_loss * static_cast<U>(gamma[l+k]) * (c_h - c_mean) * c_invvar;
+        }
+      }
+      for (;  l < n2;  ++l) {
+        const U c_h = static_cast<U>(k_input[l]);
+        const U c_loss = static_cast<U>(k_dout[l]);
+        sum_loss1 += c_loss * static_cast<U>(gamma[l]);
+        sum_loss2 += c_loss * static_cast<U>(gamma[l]) * (c_h - c_mean) * c_invvar;
+      }
+    } else {
+      int l = 4*thrx;
+      for (;  l+3 < n2;  l+=4*numx) {
+        for (int k = 0;  k < 4;  ++k) {
+          const U c_h = static_cast<U>(k_input[l+k]);
+          const U c_loss = static_cast<U>(k_dout[l+k]);
+          sum_loss1 += c_loss;
+          sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;
+        }
+      }
+      for (;  l < n2;  ++l) {
+        const U c_h = static_cast<U>(k_input[l]);
+        const U c_loss = static_cast<U>(k_dout[l]);
+        sum_loss1 += c_loss;
+        sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;
+      }
+    }
+    // intra-warp reductions
+    for (int mask = blockDim.x/2;  mask > 0;  mask /= 2) {
+      sum_loss1 += WARP_SHFL_XOR(sum_loss1, mask);
+      sum_loss2 += WARP_SHFL_XOR(sum_loss2, mask);
+    }
+    // inter-warp reductions
+    if (blockDim.y > 1) {
+      SharedMemory<U> shared;
+      U* buf = shared.getPointer(); 
+      for (int offset = blockDim.y/2;  offset > 0;  offset /= 2) {
+        // upper half of warps write to shared
+        if (threadIdx.y >= offset && threadIdx.y < 2*offset) {
+          const int wrt_i = (threadIdx.y - offset) * blockDim.x + threadIdx.x;
+          buf[2*wrt_i] = sum_loss1;
+          buf[2*wrt_i+1] = sum_loss2;
+        }
+        __syncthreads();
+        // lower half merges
+        if (threadIdx.y < offset) {
+          const int read_i = threadIdx.y * blockDim.x + threadIdx.x;
+          sum_loss1 += buf[2*read_i];
+          sum_loss2 += buf[2*read_i+1];
+        }
+        __syncthreads();
+      }
+      if (threadIdx.y == 0) {
+        buf[2*threadIdx.x] = sum_loss1;
+        buf[2*threadIdx.x+1] = sum_loss2;
+      }
+      __syncthreads();
+      if (threadIdx.y !=0) {
+        sum_loss1 = buf[2*threadIdx.x];
+        sum_loss2 = buf[2*threadIdx.x+1];
+      } 
+    }
+    // all threads now have the two sums over l
+    U fH = (U)n2;
+    U term1 = (U(1) / fH) * c_invvar;
+    T* k_grad_input = grad_input + i1*n2;
+    if (gamma != NULL) {
+      for (int l = thrx;  l < n2;  l+=numx) {
+        const U c_h = static_cast<U>(k_input[l]);
+        const U c_loss = static_cast<U>(k_dout[l]);
+        const T c_resid= static_cast<T>(k_dout_resid[l]);
+        U f_grad_input = fH * c_loss * static_cast<U>(gamma[l]);
+        f_grad_input -= sum_loss1;
+        f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;
+        f_grad_input *= term1;
+        k_grad_input[l] = static_cast<T>(f_grad_input)+c_resid;
+      }
+    } else {
+      for (int l = thrx;  l < n2;  l+=numx) {
+        const U c_h = static_cast<U>(k_input[l]);
+        const U c_loss = static_cast<U>(k_dout[l]);
+        const T c_resid= static_cast<T>(k_dout_resid[l]);
+        U f_grad_input = fH * c_loss;
+        f_grad_input -= sum_loss1;
+        f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;
+        f_grad_input *= term1;
+        k_grad_input[l] = static_cast<T>(f_grad_input)+c_resid;
+      }
+    }
+  }
+}
+
+template<typename T, typename U> 
+void HostApplyLayerNorm(
+    T* output,
+    U* mean,
+    U* invvar,
+    const T* input,
+    int n1,
+    int n2,
+    double epsilon,
+    const T* gamma,
+    const T* beta
+    )
+{
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    const dim3 threads(32,4,1);
+    const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];
+    const dim3 blocks(1, std::min((uint64_t)n1, maxGridY), 1);
+    int nshared = 
+        threads.y > 1 ? 
+	    threads.y*sizeof(U)+(threads.y/2)*sizeof(U) : 
+	    0;
+    cuApplyLayerNorm<<<blocks, threads, nshared, stream>>>(
+		    output,
+		    mean,
+		    invvar,
+		    input,
+		    n1,n2,
+		    U(epsilon),
+            gamma,beta);
+}
+
+template<typename T, typename U> 
+void HostLayerNormGradient(
+    const T* dout,
+    const T* dout_resid,
+    const U* mean,
+    const U* invvar,
+    const at::Tensor& input,
+    int n1,
+    int n2,
+    const T* gamma,
+    const T* beta,
+    double epsilon,
+    T* grad_input,
+    T* grad_gamma,
+    T* grad_beta
+    )
+{
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    if (gamma != NULL && beta != NULL) {
+      // compute grad_gamma(j) and grad_beta(j)
+      const int part_size = 16;
+      const dim3 threads2(32,4,1);
+      const dim3 blocks2((n2+threads2.x-1)/threads2.x,part_size,1);
+      const int nshared2_a = 2 * sizeof(U) * threads2.y * threads2.y * (threads2.x + 1);
+      const int nshared2_b = threads2.x * threads2.y * sizeof(U);
+      const int nshared2 = nshared2_a > nshared2_b ? nshared2_a : nshared2_b;
+      at::Tensor part_grad_gamma = at::empty({part_size,n2}, input.options().dtype(input.scalar_type()==at::ScalarType::Half ? at::ScalarType::Float : input.scalar_type()));
+      at::Tensor part_grad_beta = at::empty_like(part_grad_gamma);
+      cuComputePartGradGammaBeta<<<blocks2, threads2, nshared2, stream>>>(
+		      dout,
+		      static_cast<T*>(input.data_ptr()),
+		      n1,n2,
+		      mean,
+		      invvar,
+		      U(epsilon),
+		      static_cast<U*>(part_grad_gamma.data_ptr()),
+		      static_cast<U*>(part_grad_beta.data_ptr()));
+
+      const dim3 threads3(32,8,1);
+      const dim3 blocks3((n2+threads2.x-1)/threads2.x,1,1);
+      const int nshared3 = threads3.x * threads3.y * sizeof(U);
+      cuComputeGradGammaBeta<<<blocks3, threads3, nshared3, stream>>>(
+		      static_cast<U*>(part_grad_gamma.data_ptr()),
+		      static_cast<U*>(part_grad_beta.data_ptr()),
+		      part_size,
+		      n1,n2,
+		      grad_gamma,
+		      grad_beta);
+    }
+
+    // compute grad_input
+    const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];
+    const dim3 blocks1(1, std::min((uint64_t)n1, maxGridY), 1);
+    const dim3 threads1(32,4,1);
+    int nshared =
+	    threads1.y > 1 ?
+	    threads1.y*threads1.x*sizeof(U) :
+	    0;
+    cuComputeGradInput<<<blocks1, threads1, nshared, stream>>>(
+            dout,
+	    dout_resid,
+            static_cast<T*>(input.data_ptr()),
+            n1,n2,
+            mean,
+            invvar,
+            U(epsilon),
+            gamma,
+            grad_input);
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/masked_softmax_dropout.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/masked_softmax_dropout.cpp
new file mode 100644
index 0000000000..4d93b8a497
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/masked_softmax_dropout.cpp
@@ -0,0 +1,93 @@
+#include <torch/extension.h>
+#include <vector>
+
+namespace multihead_attn {
+namespace fused_softmax {
+namespace mask_softmax_dropout {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& input, 
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                                  );
+
+torch::Tensor bwd_cuda(
+		               int heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& dropout_mask,
+                               const uint8_t *padding_mask,
+                               float                dropout_prob
+                                                  );
+
+// C++ interface
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> fwd(
+ 			       bool 				use_mask,
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& input,
+                               torch::Tensor const& pad_mask,
+                               float                dropout_prob
+                                                 )
+{
+  AT_ASSERTM(input.dim()         == 3, "expected 3D tensor");
+  AT_ASSERTM(input.type().scalarType()         == at::ScalarType::Half, "Only HALF is supported");
+
+  if (use_mask) {
+  	AT_ASSERTM(pad_mask.dim()                     == 2,                    "expected 2D tensor");
+  	AT_ASSERTM(pad_mask.type().scalarType()       == at::ScalarType::Byte, "Only BYTE is supported");
+  }
+
+  return fwd_cuda(
+                                 is_training,
+                                 heads, 
+                                 input, 
+                                 use_mask ? static_cast<const uint8_t*>(pad_mask.data_ptr()) : nullptr, 
+                                 dropout_prob
+                                );
+}
+
+torch::Tensor bwd(
+		               bool use_mask,
+		               int heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& dropout_mask,
+                               torch::Tensor const& padding_mask,
+                               float                dropout_prob
+                                                  )
+{
+  AT_ASSERTM(output_grads.dim()      == 3, "expected 3D tensor");
+  AT_ASSERTM(softmax_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(dropout_mask.dim()      == 3, "expected 3D tensor");
+
+  AT_ASSERTM(output_grads.type().scalarType()      == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(softmax_results.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+//  AT_ASSERTM(dropout_mask.type().scalarType()      == at::ScalarType::Byte, "Only BYTE is supported");
+
+  return bwd_cuda(
+		                 heads,
+                                 output_grads,
+                                 softmax_results, 
+                                 dropout_mask, 
+                                 use_mask ? static_cast<const uint8_t*>(padding_mask.data_ptr()) : nullptr, 
+                                 dropout_prob
+                                );
+}
+
+} // end namespace mask_softmax_dropout
+} // end namespace fused_softmax
+} // end namespace multihead_attn
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &multihead_attn::fused_softmax::mask_softmax_dropout::fwd, "Self Multihead Attention masked softmax dropout -- Forward.");
+  m.def("backward", &multihead_attn::fused_softmax::mask_softmax_dropout::bwd, "Self Multihead Attention masked softmax dropout -- Backward.");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/masked_softmax_dropout_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/masked_softmax_dropout_cuda.cu
new file mode 100644
index 0000000000..318bf5814e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/masked_softmax_dropout_cuda.cu
@@ -0,0 +1,147 @@
+#include <vector>
+#include <math.h>
+#include <iostream>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include "softmax.h"
+#include "dropout.h"
+
+// symbol to be automatically resolved by PyTorch libs
+extern THCState *state;
+
+namespace multihead_attn {
+namespace fused_softmax {
+namespace mask_softmax_dropout {
+
+std::vector<torch::Tensor> fwd_cuda(
+			       bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& input, 
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   attn_batches   = input.size(0);
+  const int   sequences      = attn_batches / heads;
+  const int   q_seq_len      = input.size(1);
+  const int   k_seq_len      = q_seq_len;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+
+  // There is no reason to use more than one stream as every kernel is 
+  // sequentially dependent
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // 3 Intermediate Results + Output (Note: dropout intermediates are generated by ATen library code)
+  auto act_options  = input.options().requires_grad(false);
+  auto mask_options = act_options.dtype(torch::kUInt8);
+
+  torch::Tensor softmax_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_mask      = torch::empty({attn_batches, q_seq_len, k_seq_len},   mask_options);
+
+  // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)
+  void* input_ptr = static_cast<void*>(input.data_ptr());
+  void* softmax_results_ptr = static_cast<void*>(softmax_results.data_ptr());
+
+  // Padded Softmax
+  bool softmax_success = false;
+  if (pad_mask == nullptr) {
+    softmax_success = dispatch_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(input_ptr),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  } else {
+      softmax_success = dispatch_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(input_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             attn_batches*q_seq_len/sequences);
+  }
+
+
+  if (is_training) {
+    //use at:: function so that C++ version generates the same random mask as python version
+    auto dropout_tuple = at::_fused_dropout(softmax_results, 1.0f-dropout_prob);
+    dropout_results = std::get<0>(dropout_tuple);
+    dropout_mask = std::get<1>(dropout_tuple);
+  }
+
+  // Matmul2
+
+  return {
+           dropout_results,  
+           dropout_mask, 
+           softmax_results
+         };
+}
+
+torch::Tensor bwd_cuda(
+		               int heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& softmax_results, 
+                               torch::Tensor const& dropout_mask,
+                               const uint8_t  *padding_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   attn_batches   = output_grads.size(0);
+  const int   q_seq_len      = output_grads.size(1);
+  const int   k_seq_len      = q_seq_len;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+  // TODO: Streams can be used in Backprop but I haven't added more than one
+  // in my first attempt to create the code
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // Output Tensor Allocations
+//  torch::Tensor input_grads         = torch::empty_like(output_grads);
+
+  // Apply Dropout Mask and Scale by Dropout Probability 
+  // Softmax Grad
+  if (padding_mask == nullptr) {
+      dispatch_masked_scale_softmax_backward_stream<half, half, float,false>(
+                             static_cast<half*>(output_grads.data_ptr()), 
+                             static_cast<half*>(output_grads.data_ptr()), 
+                             reinterpret_cast<half const*>(softmax_results.data_ptr()),
+			     static_cast<uint8_t const*>(dropout_mask.data_ptr()),
+			     1.0/(1.0-dropout_prob),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len, stream);
+  } else{
+      dispatch_masked_scale_softmax_backward_masked_out_stream<half, half, float,false>(
+                             static_cast<half*>(output_grads.data_ptr()), 
+                             static_cast<half*>(output_grads.data_ptr()), 
+                             reinterpret_cast<half const*>(softmax_results.data_ptr()),
+			     static_cast<uint8_t const*>(dropout_mask.data_ptr()),
+			     static_cast<uint8_t const*>(padding_mask),
+			     1.0/(1.0-dropout_prob),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+			     heads, stream); 
+  
+  }
+//backward pass is completely in-place
+  return output_grads;
+}
+}
+}
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/philox.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/philox.h
new file mode 100644
index 0000000000..6c30a45c93
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/philox.h
@@ -0,0 +1,90 @@
+#pragma once
+//Philox CUDA. 
+
+class Philox {
+public:
+  __device__ inline Philox(unsigned long long seed,
+                           unsigned long long subsequence,
+                           unsigned long long offset) {
+    key.x = (unsigned int)seed;
+    key.y = (unsigned int)(seed >> 32);
+    counter = make_uint4(0, 0, 0, 0);
+    counter.z = (unsigned int)(subsequence);
+    counter.w = (unsigned int)(subsequence >> 32);
+    STATE = 0;
+    incr_n(offset / 4);
+  }
+  __device__ inline uint4 operator()() {
+    if(STATE == 0) {
+      uint4 counter_ = counter;
+      uint2 key_ = key;
+      //7-round philox
+      for(int i = 0; i < 6; i++) {
+        counter_ = single_round(counter_, key_);
+        key_.x += (kPhilox10A); key_.y += (kPhilox10B);
+      }
+      output = single_round(counter_, key_);
+      incr();
+    }
+    //return a float4 directly
+    //unsigned long ret;
+    //switch(STATE) {
+    //  case 0: ret = output.x; break;
+    //  case 1: ret = output.y; break;
+    //  case 2: ret = output.z; break;
+    //  case 3: ret = output.w; break;
+    //}
+    //STATE = (STATE + 1) % 4;
+    return output;
+  }
+private:
+  uint4 counter;
+  uint4 output;
+  uint2 key;
+  unsigned int STATE;
+  __device__ inline void incr_n(unsigned long long n) {
+    unsigned int nlo = (unsigned int)(n);
+    unsigned int nhi = (unsigned int)(n >> 32);
+    counter.x += nlo;
+    if (counter.x < nlo)
+      nhi++;
+    counter.y += nhi;
+    if (nhi <= counter.y)
+      return;
+    if (++counter.z)
+      return;
+    ++counter.w;
+  }
+  __device__ inline void incr() {
+    if (++counter.x)
+      return;
+    if (++counter.y)
+      return;
+    if (++counter.z)
+      return;
+    ++counter.w;
+  }
+  __device__ unsigned int mulhilo32(unsigned int a, unsigned int b,
+                                    unsigned int *result_high) {
+    *result_high = __umulhi(a, b);
+    return a*b;
+  }
+  __device__ inline uint4 single_round(uint4 ctr, uint2 key) {
+    unsigned int hi0;
+    unsigned int hi1;
+    unsigned int lo0 = mulhilo32(kPhiloxSA, ctr.x, &hi0);
+    unsigned int lo1 = mulhilo32(kPhiloxSB, ctr.z, &hi1);
+    uint4 ret = {hi1 ^ ctr.y ^ key.x, lo1, hi0 ^ ctr.w ^ key.y, lo0};
+    return ret;
+  }
+  static const unsigned long kPhilox10A = 0x9E3779B9;
+  static const unsigned long kPhilox10B = 0xBB67AE85;
+  static const unsigned long kPhiloxSA = 0xD2511F53;
+  static const unsigned long kPhiloxSB = 0xCD9E8D57;
+};
+// Inverse of 2^32.
+#define M_RAN_INVM32 2.3283064e-10f
+__device__  __inline__ float4 uniform4(uint4 x) {
+    return make_float4(x.x * M_RAN_INVM32, x.y * M_RAN_INVM32, x.z * M_RAN_INVM32,x.w * M_RAN_INVM32);
+
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn.cpp
new file mode 100644
index 0000000000..ba096516db
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn.cpp
@@ -0,0 +1,132 @@
+#include <torch/extension.h>
+#include <vector>
+
+namespace multihead_attn {
+namespace self {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,  
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                                  );
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  );
+
+// C++ interface
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> fwd(
+ 							   bool 				use_mask,
+                               bool                 use_time_mask,
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& pad_mask,
+                               float                dropout_prob
+                                                 )
+{
+  AT_ASSERTM(inputs.dim()         == 3, "expected 3D tensor");
+  AT_ASSERTM(input_weights.dim()  == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim() == 2, "expected 2D tensor");
+
+  AT_ASSERTM(inputs.type().scalarType()         == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights.type().scalarType()  == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType() == at::ScalarType::Half, "Only HALF is supported");
+
+  if (use_mask) {
+  	AT_ASSERTM(pad_mask.dim()                     == 2,                    "expected 2D tensor");
+  	AT_ASSERTM(pad_mask.type().scalarType()       == at::ScalarType::Byte, "Only BYTE is supported");
+  }
+  
+  return fwd_cuda(
+                                 use_time_mask,
+                                 is_training,
+                                 heads, 
+                                 inputs, 
+                                 input_weights, 
+                                 output_weights, 
+                                 use_mask ? static_cast<const uint8_t*>(pad_mask.data_ptr()) : nullptr, 
+                                 dropout_prob
+                                );
+}
+
+std::vector<torch::Tensor> bwd(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  )
+{
+  AT_ASSERTM(output_grads.dim()      == 3, "expected 3D tensor");
+  AT_ASSERTM(matmul2_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(dropout_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(softmax_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(input_lin_results.dim() == 3, "expected 3D tensor");
+  AT_ASSERTM(inputs.dim()            == 3, "expected 3D tensor");
+  AT_ASSERTM(input_weights.dim()     == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim()    == 2, "expected 2D tensor");
+  AT_ASSERTM(dropout_mask.dim()      == 3, "expected 3D tensor");
+  
+  AT_ASSERTM(output_grads.type().scalarType()      == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(matmul2_results.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(dropout_results.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(softmax_results.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_lin_results.type().scalarType() == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(inputs.type().scalarType()            == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights.type().scalarType()     == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType()    == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(dropout_mask.type().scalarType()      == at::ScalarType::Byte, "Only BYTE is supported");
+  
+  return bwd_cuda(
+                                 heads, 
+                                 output_grads,
+                                 matmul2_results,
+                                 dropout_results,
+                                 softmax_results, 
+                                 input_lin_results, 
+                                 inputs, 
+                                 input_weights,
+                                 output_weights,
+                                 dropout_mask, 
+                                 dropout_prob
+                                );
+}
+
+} // end namespace cublas_gemmex
+} // end namespace self
+} // end namespace multihead_attn
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &multihead_attn::self::cublas_gemmex::fwd, "Self Multihead Attention Forward.");
+  m.def("backward", &multihead_attn::self::cublas_gemmex::bwd, "Self Multihead Attention Backward.");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias.cpp
new file mode 100644
index 0000000000..9ed393a40d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias.cpp
@@ -0,0 +1,139 @@
+#include <torch/extension.h>
+#include <vector>
+
+namespace multihead_attn {
+namespace self_bias {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,  
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& input_biases,
+                               torch::Tensor const& output_biases,
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                                  );
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               //torch::Tensor const& input_biases,
+                               //torch::Tensor const& output_biases,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  );
+
+// C++ interface
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> fwd(
+ 							   bool 				use_mask,
+                               bool                 use_time_mask,
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& input_biases, torch::Tensor const& output_biases,
+                               torch::Tensor const& pad_mask,
+                               float                dropout_prob
+                                                 )
+{
+  AT_ASSERTM(inputs.dim()         == 3, "expected 3D tensor");
+  AT_ASSERTM(input_weights.dim()  == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim() == 2, "expected 2D tensor");
+
+  AT_ASSERTM(inputs.type().scalarType()         == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights.type().scalarType()  == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType() == at::ScalarType::Half, "Only HALF is supported");
+
+  if (use_mask) {
+  	AT_ASSERTM(pad_mask.dim()                     == 2,                    "expected 2D tensor");
+  	AT_ASSERTM(pad_mask.type().scalarType()       == at::ScalarType::Byte, "Only BYTE is supported");
+  }
+
+  return fwd_cuda(
+                                 use_time_mask,
+                                 is_training,
+                                 heads, 
+                                 inputs, 
+                                 input_weights, 
+                                 output_weights, 
+                                 input_biases, 
+                                 output_biases, 
+                                 use_mask ? static_cast<const uint8_t*>(pad_mask.data_ptr()) : nullptr, 
+                                 dropout_prob
+                                );
+}
+
+std::vector<torch::Tensor> bwd(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  )
+{
+  AT_ASSERTM(output_grads.dim()      == 3, "expected 3D tensor");
+  AT_ASSERTM(matmul2_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(dropout_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(softmax_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(input_lin_results.dim() == 3, "expected 3D tensor");
+  AT_ASSERTM(inputs.dim()            == 3, "expected 3D tensor");
+  AT_ASSERTM(input_weights.dim()     == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim()    == 2, "expected 2D tensor");
+  AT_ASSERTM(dropout_mask.dim()      == 3, "expected 3D tensor");
+
+  AT_ASSERTM(output_grads.type().scalarType()      == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(matmul2_results.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(dropout_results.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(softmax_results.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_lin_results.type().scalarType() == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(inputs.type().scalarType()            == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights.type().scalarType()     == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType()    == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(dropout_mask.type().scalarType()      == at::ScalarType::Byte, "Only BYTE is supported");
+
+  return bwd_cuda(
+                                 heads, 
+                                 output_grads,
+                                 matmul2_results,
+                                 dropout_results,
+                                 softmax_results, 
+                                 input_lin_results, 
+                                 inputs, 
+                                 input_weights,
+                                 output_weights,
+                                 dropout_mask, 
+                                 dropout_prob
+                                );
+}
+
+} // end namespace cublas_gemmex
+} // end namespace self
+} // end namespace multihead_attn
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &multihead_attn::self_bias::cublas_gemmex::fwd, "Self Multihead Attention with Bias -- Forward.");
+  m.def("backward", &multihead_attn::self_bias::cublas_gemmex::bwd, "Self Multihead Attention with Bias -- Backward.");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask.cpp
new file mode 100644
index 0000000000..2bca0ade0d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask.cpp
@@ -0,0 +1,143 @@
+#include <torch/extension.h>
+#include <vector>
+#include <cuda_fp16.h>
+
+namespace multihead_attn {
+namespace self_bias_additive_mask {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,  
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& input_biases,
+                               torch::Tensor const& output_biases,
+                               const half*       pad_mask,
+                               float                dropout_prob
+                                                  );
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                              // torch::Tensor const& softmax_results,
+                               torch::Tensor const& bmm1_results,
+                               torch::Tensor const& pad_mask,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               //torch::Tensor const& input_biases,
+                               //torch::Tensor const& output_biases,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  );
+
+// C++ interface
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> fwd(
+ 							   bool 				use_mask,
+                               bool                 use_time_mask,
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& input_biases, torch::Tensor const& output_biases,
+                               torch::Tensor const& pad_mask,
+                               float                dropout_prob
+                                                 )
+{
+  AT_ASSERTM(inputs.dim()         == 3, "expected 3D tensor");
+  AT_ASSERTM(input_weights.dim()  == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim() == 2, "expected 2D tensor");
+
+  AT_ASSERTM(inputs.type().scalarType()         == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights.type().scalarType()  == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType() == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(use_mask                                                  , "no mask is not supported");
+
+  if (use_mask) {
+  	AT_ASSERTM(pad_mask.dim()                     == 2,                    "expected 2D tensor");
+  	AT_ASSERTM(pad_mask.type().scalarType()       == at::ScalarType::Half, "Only Half is supported");
+  }
+
+  return fwd_cuda(
+                                 use_time_mask,
+                                 is_training,
+                                 heads, 
+                                 inputs, 
+                                 input_weights, 
+                                 output_weights, 
+                                 input_biases, 
+                                 output_biases, 
+                                 use_mask ? static_cast<const half*>(pad_mask.data_ptr()) : nullptr, 
+                                 dropout_prob
+                                );
+}
+
+std::vector<torch::Tensor> bwd(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& bmm1_results,
+                               torch::Tensor const& pad_mask,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                                  )
+{
+  AT_ASSERTM(output_grads.dim()      == 3, "expected 3D tensor");
+  AT_ASSERTM(matmul2_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(dropout_results.dim()   == 3, "expected 3D tensor");
+  AT_ASSERTM(input_lin_results.dim() == 3, "expected 3D tensor");
+  AT_ASSERTM(inputs.dim()            == 3, "expected 3D tensor");
+  AT_ASSERTM(input_weights.dim()     == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim()    == 2, "expected 2D tensor");
+  AT_ASSERTM(dropout_mask.dim()      == 3, "expected 3D tensor");
+
+  AT_ASSERTM(output_grads.type().scalarType()      == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(matmul2_results.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(dropout_results.type().scalarType()   == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_lin_results.type().scalarType() == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(inputs.type().scalarType()            == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights.type().scalarType()     == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType()    == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(dropout_mask.type().scalarType()      == at::ScalarType::Byte, "Only BYTE is supported");
+
+  return bwd_cuda(
+                                 heads, 
+                                 output_grads,
+                                 matmul2_results,
+                                 dropout_results,
+				 bmm1_results,
+				 pad_mask, 
+                                 input_lin_results, 
+                                 inputs, 
+                                 input_weights,
+                                 output_weights,
+                                 dropout_mask, 
+                                 dropout_prob
+                                );
+}
+
+} // end namespace cublas_gemmex
+} // end namespace self
+} // end namespace multihead_attn
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &multihead_attn::self_bias_additive_mask::cublas_gemmex::fwd, "Self Multihead Attention with Bias -- Forward.");
+  m.def("backward", &multihead_attn::self_bias_additive_mask::cublas_gemmex::bwd, "Self Multihead Attention with Bias -- Backward.");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask_cuda.cu
new file mode 100644
index 0000000000..3dea68a2d2
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask_cuda.cu
@@ -0,0 +1,463 @@
+#include <vector>
+#include <math.h>
+#include <iostream>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include "strided_batched_gemm.h"
+#include "softmax.h"
+#include "dropout.h"
+#include "layer_norm.h"
+
+// symbol to be automatically resolved by PyTorch libs
+extern THCState *state;
+
+namespace multihead_attn {
+namespace self_bias_additive_mask {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,
+							   bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& input_biases,
+                               torch::Tensor const& output_biases,
+                               const half*       pad_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim      = inputs.size(2);
+  const int   sequences      = inputs.size(1);
+  const int   q_seq_len      = inputs.size(0);
+  const int   k_seq_len      = q_seq_len;
+  const int   batches        = sequences * q_seq_len;
+  const int   head_dim       = embed_dim / heads;
+  const int   output_lin_dim = 3 * embed_dim;
+  const int   attn_batches   = heads * sequences;
+  const int   lead_dim       = attn_batches * 3 * head_dim;
+  const int   batch_stride   = 3 * head_dim;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+  const float alpha          = 1.0;
+  const float beta_zero       = 0.0;
+  const float beta_one           = 1.0;
+  const float scale          = 1.0 / sqrt(static_cast<float>(head_dim));
+
+  // There is no reason to use more than one stream as every kernel is 
+  // sequentially dependent
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // 3 Intermediate Results + Output (Note: dropout intermediates are generated by ATen library code)
+  auto act_options  = inputs.options().requires_grad(false);
+  auto mask_options = act_options.dtype(torch::kUInt8);
+
+  torch::Tensor input_lin_results = torch::empty({q_seq_len, sequences, output_lin_dim}, act_options);
+  torch::Tensor bmm1_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_mask      = torch::empty({attn_batches, q_seq_len, k_seq_len},   mask_options);
+  torch::Tensor matmul2_results   = torch::empty({q_seq_len, attn_batches, head_dim},    act_options);
+  torch::Tensor outputs           = torch::empty_like(inputs, act_options);
+
+  // Input Linear Results Pointers to Q, K, and V of interviewed activations
+  void* q_lin_results_ptr   = static_cast<void*>(input_lin_results.data_ptr());
+  void* k_lin_results_ptr   = static_cast<void*>(static_cast<half*>(input_lin_results.data_ptr()) + head_dim);
+  void* v_lin_results_ptr   = static_cast<void*>(static_cast<half*>(input_lin_results.data_ptr()) + 2*head_dim);
+
+  // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)
+  void* bmm1_results_ptr = static_cast<void*>(bmm1_results.data_ptr());
+  void* dropout_results_ptr = static_cast<void*>(dropout_results.data_ptr());
+
+  char a_layout_t{'t'};
+  char a_layout_n{'n'};
+  char b_layout_n{'n'};
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+  // Input Linear Fwd
+  input_lin_results.copy_(input_biases);
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             output_lin_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(inputs.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta_one),
+                             q_lin_results_ptr,
+                             CUDA_R_16F, 
+                             output_lin_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             scale, 
+                             static_cast<const half*>(k_lin_results_ptr), 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<const half*>(q_lin_results_ptr),
+                             lead_dim, 
+                             batch_stride, 
+                             beta_zero, 
+                             static_cast<half*>(bmm1_results_ptr), 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             attn_batches);
+  // Padded Softmax
+  bool softmax_success = false;
+  if (is_training) {
+      softmax_success = dispatch_additive_masked_softmax_dropout<half, half, float>(
+                           reinterpret_cast<half*>(dropout_results_ptr),
+                           (is_training) ? reinterpret_cast<uint8_t*>(dropout_mask.data_ptr<uint8_t>()) : nullptr,
+                           reinterpret_cast<const half*>(bmm1_results_ptr),
+                           pad_mask,
+      		           attn_batches*q_seq_len*q_seq_len,
+                           k_seq_len,
+                           k_seq_len,
+                           attn_batches*q_seq_len,
+                           attn_batches*q_seq_len/sequences, 
+      		           1.0f-dropout_prob,
+		           stream);
+  } else {
+      softmax_success = dispatch_additive_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(dropout_results_ptr),//this is actually softmax results, but making it consistent for the next function
+                             reinterpret_cast<const half*>(bmm1_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             attn_batches*q_seq_len/sequences);
+  }
+
+  // Matmul2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr), 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<const half*>(dropout_results.data_ptr()), 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta_zero, 
+                             static_cast<half*>(matmul2_results.data_ptr()), 
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             attn_batches);
+
+  outputs.copy_(output_biases);
+
+  // Output Linear
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta_one),
+                             static_cast<void*>(outputs.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO1_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return {
+           input_lin_results,  
+           bmm1_results,
+           dropout_results, 
+           dropout_mask, 
+           matmul2_results, 
+           outputs
+         };
+}
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& bmm1_results,
+                               torch::Tensor const& pad_mask,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim      = inputs.size(2);
+  const int   sequences      = inputs.size(1);
+  const int   q_seq_len      = inputs.size(0);
+  const int   k_seq_len      = q_seq_len;
+  const int   batches        = sequences * q_seq_len;
+  const int   head_dim       = embed_dim / heads;
+  const int   output_lin_dim = 3 * embed_dim;
+  const int   attn_batches   = heads * sequences;
+  const int   lead_dim       = attn_batches * 3 * head_dim;
+  const int   batch_stride   = 3 * head_dim;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+  const float alpha          = 1.0;
+  const float beta           = 0.0;
+  const float scale          = 1.0 / sqrt(static_cast<float>(head_dim));
+
+  // TODO: Streams can be used in Backprop but I haven't added more than one
+  // in my first attempt to create the code
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // Output Tensor Allocations
+  torch::Tensor input_grads         = torch::empty_like(inputs);
+  torch::Tensor input_weight_grads  = torch::empty_like(input_weights);
+  torch::Tensor output_weight_grads = torch::empty_like(output_weights);
+  // Intermediate Tensor Allocations
+  at::Tensor output_lin_grads       = torch::empty_like(matmul2_results);
+  at::Tensor matmul2_grads          = torch::empty_like(dropout_results);
+  at::Tensor input_lin_output_grads = torch::empty_like(input_lin_results);
+
+  auto q_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr());
+  auto k_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr()) + head_dim;
+  auto v_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr()) + 2*head_dim;
+
+  auto q_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr());
+  auto k_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr()) + head_dim;
+  auto v_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr()) + 2*head_dim;
+
+  char a_layout_n{'n'};
+  char a_layout_t{'t'};
+  char b_layout_n{'n'};
+  char b_layout_t{'t'}; 
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+
+  // Output Linear Dgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(output_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_lin_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  // Output Linear Wgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             embed_dim,
+                             batches, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(output_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_weight_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  auto  output_bias_grads = output_grads.view({-1, embed_dim}) .sum(0, false);
+  // MatMul2 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr),
+                             lead_dim, 
+                             batch_stride,
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             beta, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len,
+                             attn_batches);
+
+  // Matmul2 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             static_cast<const half*>(dropout_results.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             v_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+
+  // Apply Dropout Mask and Scale by Dropout Probability 
+  // Softmax Grad
+  dispatch_masked_scale_softmax_backward_recompute<half, half, float, false>(
+                             static_cast<half*>(matmul2_grads.data_ptr()), 
+                             static_cast<half* const>(matmul2_grads.data_ptr()), 
+                             reinterpret_cast<half const*>(bmm1_results.data_ptr()),
+                             reinterpret_cast<half const*>(pad_mask.data_ptr()),
+			     static_cast<uint8_t const*>(dropout_mask.data_ptr()),
+			     1.0/(1.0-dropout_prob),
+                             k_seq_len,
+                             k_seq_len,
+			     attn_batches*q_seq_len/sequences,
+                             attn_batches*q_seq_len,
+			     stream);
+
+  // Matmul1 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             scale, 
+                             k_lin_results_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             q_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+
+  // Matmul1 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             scale, 
+                             q_lin_results_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             k_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+  // Input Linear Dgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim,
+                             batches, 
+                             output_lin_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+			     static_cast<const void*>(input_lin_output_grads.data_ptr()),
+                             //static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F, 
+                             output_lin_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO10_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  // Input Linear Wgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             output_lin_dim,
+                             batches, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(inputs.data_ptr()),
+                             CUDA_R_16F,
+                             embed_dim,
+                             static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F,
+                             output_lin_dim,
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_weight_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  auto  input_bias_grads = input_lin_output_grads.view({-1, output_lin_dim}).sum(0, false);
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return { 
+           input_grads, 
+           input_weight_grads, 
+           output_weight_grads,
+           input_bias_grads, 
+           output_bias_grads
+         };
+}
+
+} // end namespace cublas_gemmex
+} // end namespace self
+} // end namespace multihead_attn
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_cuda.cu
new file mode 100644
index 0000000000..fd33c9757b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_cuda.cu
@@ -0,0 +1,471 @@
+#include <vector>
+#include <math.h>
+#include <iostream>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include "strided_batched_gemm.h"
+#include "softmax.h"
+#include "dropout.h"
+#include "layer_norm.h"
+
+// symbol to be automatically resolved by PyTorch libs
+extern THCState *state;
+
+namespace multihead_attn {
+namespace self_bias {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,
+							   bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& input_biases,
+                               torch::Tensor const& output_biases,
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim      = inputs.size(2);
+  const int   sequences      = inputs.size(1);
+  const int   q_seq_len      = inputs.size(0);
+  const int   k_seq_len      = q_seq_len;
+  const int   batches        = sequences * q_seq_len;
+  const int   head_dim       = embed_dim / heads;
+  const int   output_lin_dim = 3 * embed_dim;
+  const int   attn_batches   = heads * sequences;
+  const int   lead_dim       = attn_batches * 3 * head_dim;
+  const int   batch_stride   = 3 * head_dim;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+  const float alpha          = 1.0;
+  const float beta_zero       = 0.0;
+  const float beta_one           = 1.0;
+  const float scale          = 1.0 / sqrt(static_cast<float>(head_dim));
+
+  // There is no reason to use more than one stream as every kernel is 
+  // sequentially dependent
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // 3 Intermediate Results + Output (Note: dropout intermediates are generated by ATen library code)
+  auto act_options  = inputs.options().requires_grad(false);
+  auto mask_options = act_options.dtype(torch::kUInt8);
+
+  torch::Tensor input_lin_results = torch::empty({q_seq_len, sequences, output_lin_dim}, act_options);
+  torch::Tensor softmax_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_mask      = torch::empty({attn_batches, q_seq_len, k_seq_len},   mask_options);
+  torch::Tensor matmul2_results   = torch::empty({q_seq_len, attn_batches, head_dim},    act_options);
+  torch::Tensor outputs           = torch::empty_like(inputs, act_options);
+
+  // Input Linear Results Pointers to Q, K, and V of interviewed activations
+  void* q_lin_results_ptr   = static_cast<void*>(input_lin_results.data_ptr());
+  void* k_lin_results_ptr   = static_cast<void*>(static_cast<half*>(input_lin_results.data_ptr()) + head_dim);
+  void* v_lin_results_ptr   = static_cast<void*>(static_cast<half*>(input_lin_results.data_ptr()) + 2*head_dim);
+
+  // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)
+  void* softmax_results_ptr = static_cast<void*>(softmax_results.data_ptr());
+
+  char a_layout_t{'t'};
+  char a_layout_n{'n'};
+  char b_layout_n{'n'};
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+  // Input Linear Fwd
+  input_lin_results.copy_(input_biases);
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             output_lin_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(inputs.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta_one),
+                             q_lin_results_ptr,
+                             CUDA_R_16F, 
+                             output_lin_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             scale, 
+                             static_cast<const half*>(k_lin_results_ptr), 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<const half*>(q_lin_results_ptr),
+                             lead_dim, 
+                             batch_stride, 
+                             beta_zero, 
+                             static_cast<half*>(softmax_results_ptr), 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             attn_batches);
+  // Padded Softmax
+  bool softmax_success = false;
+  if (pad_mask == nullptr) {
+    softmax_success = dispatch_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  } else {
+    if (use_time_mask) {
+      softmax_success = dispatch_time_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             q_seq_len);
+    } else {
+      softmax_success = dispatch_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             attn_batches*q_seq_len/sequences);
+    }
+  }
+
+
+  if (is_training) {
+    //use at:: function so that C++ version generates the same random mask as python version
+    auto dropout_tuple = at::_fused_dropout(softmax_results, 1.0f-dropout_prob);
+    dropout_results = std::get<0>(dropout_tuple);
+    dropout_mask = std::get<1>(dropout_tuple);
+  }
+
+  // Matmul2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr), 
+                             lead_dim, 
+                             batch_stride, 
+                             (is_training) ? static_cast<const half*>(dropout_results.data_ptr()) : static_cast<const half*>(softmax_results.data_ptr()) , 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta_zero, 
+                             static_cast<half*>(matmul2_results.data_ptr()), 
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             attn_batches);
+
+  outputs.copy_(output_biases);
+
+  // Output Linear
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta_one),
+                             static_cast<void*>(outputs.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO1_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return {
+           input_lin_results,  
+           softmax_results,
+           dropout_results, 
+           dropout_mask, 
+           matmul2_results, 
+           outputs
+         };
+}
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim      = inputs.size(2);
+  const int   sequences      = inputs.size(1);
+  const int   q_seq_len      = inputs.size(0);
+  const int   k_seq_len      = q_seq_len;
+  const int   batches        = sequences * q_seq_len;
+  const int   head_dim       = embed_dim / heads;
+  const int   output_lin_dim = 3 * embed_dim;
+  const int   attn_batches   = heads * sequences;
+  const int   lead_dim       = attn_batches * 3 * head_dim;
+  const int   batch_stride   = 3 * head_dim;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+  const float alpha          = 1.0;
+  const float beta           = 0.0;
+  const float scale          = 1.0 / sqrt(static_cast<float>(head_dim));
+
+  // TODO: Streams can be used in Backprop but I haven't added more than one
+  // in my first attempt to create the code
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // Output Tensor Allocations
+  torch::Tensor input_grads         = torch::empty_like(inputs);
+  torch::Tensor input_weight_grads  = torch::empty_like(input_weights);
+  torch::Tensor output_weight_grads = torch::empty_like(output_weights);
+  // Intermediate Tensor Allocations
+  at::Tensor output_lin_grads       = torch::empty_like(matmul2_results);
+  at::Tensor matmul2_grads          = torch::empty_like(dropout_results);
+  at::Tensor input_lin_output_grads = torch::empty_like(input_lin_results);
+
+  auto q_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr());
+  auto k_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr()) + head_dim;
+  auto v_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr()) + 2*head_dim;
+
+  auto q_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr());
+  auto k_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr()) + head_dim;
+  auto v_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr()) + 2*head_dim;
+
+  char a_layout_n{'n'};
+  char a_layout_t{'t'};
+  char b_layout_n{'n'};
+  char b_layout_t{'t'}; 
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+
+  // Output Linear Dgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(output_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_lin_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  // Output Linear Wgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             embed_dim,
+                             batches, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(output_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_weight_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  auto  output_bias_grads = output_grads.view({-1, embed_dim}) .sum(0, false);
+  // MatMul2 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr),
+                             lead_dim, 
+                             batch_stride,
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             beta, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len,
+                             attn_batches);
+
+  // Matmul2 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             static_cast<const half*>(dropout_results.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             v_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+
+  // Apply Dropout Mask and Scale by Dropout Probability 
+  // Softmax Grad
+  dispatch_masked_scale_softmax_backward_stream<half, half, float,false>(
+                             static_cast<half*>(matmul2_grads.data_ptr()), 
+                             static_cast<half*>(matmul2_grads.data_ptr()), 
+                             reinterpret_cast<half const*>(softmax_results.data_ptr()),
+			     static_cast<uint8_t const*>(dropout_mask.data_ptr()),
+			     1.0/(1.0-dropout_prob),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len, stream);
+
+  // Matmul1 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             scale, 
+                             k_lin_results_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             q_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+
+  // Matmul1 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             scale, 
+                             q_lin_results_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             k_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+  // Input Linear Dgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim,
+                             batches, 
+                             output_lin_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+			     static_cast<const void*>(input_lin_output_grads.data_ptr()),
+                             //static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F, 
+                             output_lin_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO10_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  // Input Linear Wgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             output_lin_dim,
+                             batches, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(inputs.data_ptr()),
+                             CUDA_R_16F,
+                             embed_dim,
+                             static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F,
+                             output_lin_dim,
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_weight_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  auto  input_bias_grads = input_lin_output_grads.view({-1, output_lin_dim}).sum(0, false);
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return { 
+           input_grads, 
+           input_weight_grads, 
+           output_weight_grads,
+           input_bias_grads, 
+           output_bias_grads
+         };
+}
+
+} // end namespace cublas_gemmex
+} // end namespace self
+} // end namespace multihead_attn
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_cuda.cu
new file mode 100644
index 0000000000..68d0a43d39
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_cuda.cu
@@ -0,0 +1,469 @@
+#include <vector>
+#include <math.h>
+#include <iostream>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include "strided_batched_gemm.h"
+#include "softmax.h"
+#include "dropout.h"
+#include "layer_norm.h"
+
+// symbol to be automatically resolved by PyTorch libs
+extern THCState *state;
+
+namespace multihead_attn {
+namespace self {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,
+							   bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim      = inputs.size(2);
+  const int   sequences      = inputs.size(1);
+  const int   q_seq_len      = inputs.size(0);
+  const int   k_seq_len      = q_seq_len;
+  const int   batches        = sequences * q_seq_len;
+  const int   head_dim       = embed_dim / heads;
+  const int   output_lin_dim = 3 * embed_dim;
+  const int   attn_batches   = heads * sequences;
+  const int   lead_dim       = attn_batches * 3 * head_dim;
+  const int   batch_stride   = 3 * head_dim;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+  const float alpha          = 1.0;
+  const float beta           = 0.0;
+  const float scale          = 1.0 / sqrt(static_cast<float>(head_dim));
+ 
+  // There is no reason to use more than one stream as every kernel is 
+  // sequentially dependent
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // 3 Intermediate Results + Output (Note: dropout intermediates are generated by ATen library code)
+  auto act_options  = inputs.options().requires_grad(false);
+  auto mask_options = act_options.dtype(torch::kUInt8);
+
+  torch::Tensor input_lin_results = torch::empty({q_seq_len, sequences, output_lin_dim}, act_options);
+  torch::Tensor softmax_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_mask      = torch::empty({attn_batches, q_seq_len, k_seq_len},   mask_options);
+  torch::Tensor matmul2_results   = torch::empty({q_seq_len, attn_batches, head_dim},    act_options);
+  torch::Tensor outputs           = torch::empty_like(inputs, act_options);
+
+  // Input Linear Results Pointers to Q, K, and V of interviewed activations
+  void* q_lin_results_ptr   = static_cast<void*>(input_lin_results.data_ptr());
+  void* k_lin_results_ptr   = static_cast<void*>(static_cast<half*>(input_lin_results.data_ptr()) + head_dim);
+  void* v_lin_results_ptr   = static_cast<void*>(static_cast<half*>(input_lin_results.data_ptr()) + 2*head_dim);
+
+  // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)
+  void* softmax_results_ptr = static_cast<void*>(softmax_results.data_ptr());
+  
+  char a_layout_t{'t'};
+  char a_layout_n{'n'};
+  char b_layout_n{'n'};
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+  // Input Linear Fwd
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             output_lin_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(inputs.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             q_lin_results_ptr,
+                             CUDA_R_16F, 
+                             output_lin_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             scale, 
+                             static_cast<const half*>(k_lin_results_ptr), 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<const half*>(q_lin_results_ptr),
+                             lead_dim, 
+                             batch_stride, 
+                             beta, 
+                             static_cast<half*>(softmax_results_ptr), 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             attn_batches);
+
+  // Padded Softmax
+  bool softmax_success = false;
+  if (pad_mask == nullptr) {
+    softmax_success = dispatch_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  } else {
+    if (use_time_mask) {
+      softmax_success = dispatch_time_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             q_seq_len);
+    } else {
+      softmax_success = dispatch_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             attn_batches*q_seq_len/sequences);
+    }
+  }
+  assert(softmax_success);
+
+  if (is_training) {
+    apex_fused_dropout_cuda<at::Half,float,uint32_t>(
+                               static_cast<at::Half const*>(softmax_results.data_ptr()),
+                               static_cast<at::Half*>(dropout_results.data_ptr()),
+                               static_cast<uint8_t*>(dropout_mask.data_ptr()),
+                               dropout_elems,
+                               (1.0f - dropout_prob));
+  }
+ 
+  // Matmul2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr), 
+                             lead_dim, 
+                             batch_stride, 
+                             (is_training) ? static_cast<const half*>(dropout_results.data_ptr()) : static_cast<const half*>(softmax_results.data_ptr()) , 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             static_cast<half*>(matmul2_results.data_ptr()), 
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             attn_batches);
+
+  // Output Linear
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(outputs.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return { 
+           input_lin_results, 
+           softmax_results,
+           dropout_results, 
+           dropout_mask, 
+           matmul2_results, 
+           outputs
+         };
+}
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim      = inputs.size(2);
+  const int   sequences      = inputs.size(1);
+  const int   q_seq_len      = inputs.size(0);
+  const int   k_seq_len      = q_seq_len;
+  const int   batches        = sequences * q_seq_len;
+  const int   head_dim       = embed_dim / heads;
+  const int   output_lin_dim = 3 * embed_dim;
+  const int   attn_batches   = heads * sequences;
+  const int   lead_dim       = attn_batches * 3 * head_dim;
+  const int   batch_stride   = 3 * head_dim;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+  const float alpha          = 1.0;
+  const float beta           = 0.0;
+  const float scale          = 1.0 / sqrt(static_cast<float>(head_dim));
+
+  // TODO: Streams can be used in Backprop but I haven't added more than one
+  // in my first attempt to create the code
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+ 
+  // Output Tensor Allocations
+  torch::Tensor input_grads         = torch::empty_like(inputs);
+  torch::Tensor input_weight_grads  = torch::empty_like(input_weights);
+  torch::Tensor output_weight_grads = torch::empty_like(output_weights);
+  // Intermediate Tensor Allocations
+  at::Tensor output_lin_grads       = torch::empty_like(matmul2_results);
+  at::Tensor matmul2_grads          = torch::empty_like(dropout_results);
+  at::Tensor input_lin_output_grads = torch::empty_like(input_lin_results);
+ 
+  auto q_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr());
+  auto k_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr()) + head_dim;
+  auto v_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr()) + 2*head_dim;
+  
+  auto q_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr());
+  auto k_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr()) + head_dim;
+  auto v_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr()) + 2*head_dim;
+
+  char a_layout_n{'n'};
+  char a_layout_t{'t'};
+  char b_layout_n{'n'};
+  char b_layout_t{'t'}; 
+  
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+ 
+  // Output Linear Dgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(output_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_lin_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+ 
+  // Output Linear Wgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             embed_dim,
+                             batches, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(output_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_weight_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // MatMul2 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr),
+                             lead_dim, 
+                             batch_stride,
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             beta, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len,
+                             attn_batches);
+  
+  // Matmul2 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             static_cast<const half*>(dropout_results.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             v_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+
+  // Apply Dropout Mask and Scale by Dropout Probability 
+  apex_masked_scale_cuda<at::Half,float,uint32_t>(
+                             static_cast<at::Half const*>(matmul2_grads.data_ptr()),
+                             static_cast<at::Half*>(matmul2_grads.data_ptr()),
+                             static_cast<uint8_t const*>(dropout_mask.data_ptr()),
+                             dropout_elems,
+                             (1.0 / (1.0 - dropout_prob)));
+
+  // Softmax Grad
+  bool softmax_success = false;
+  softmax_success = dispatch_softmax_backward<half, half, float>(
+                             static_cast<half*>(matmul2_grads.data_ptr()), 
+                             static_cast<half*>(matmul2_grads.data_ptr()), 
+                             reinterpret_cast<half const*>(softmax_results.data_ptr()),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  assert(softmax_success);
+
+  // Matmul1 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             scale, 
+                             k_lin_results_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             q_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+  
+  // Matmul1 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             scale, 
+                             q_lin_results_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             k_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+
+  // Input Linear Dgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim,
+                             batches, 
+                             output_lin_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F, 
+                             output_lin_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // Input Linear Wgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             output_lin_dim,
+                             batches, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(inputs.data_ptr()),
+                             CUDA_R_16F,
+                             embed_dim,
+                             static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F,
+                             output_lin_dim,
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_weight_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return { 
+           input_grads, 
+           input_weight_grads, 
+           output_weight_grads
+         };
+}
+
+} // end namespace cublas_gemmex
+} // end namespace self
+} // end namespace multihead_attn
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add.cpp
new file mode 100644
index 0000000000..a5442bc300
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add.cpp
@@ -0,0 +1,173 @@
+#include <torch/extension.h>
+#include <vector>
+
+namespace multihead_attn {
+namespace self_norm_add {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, 
+							   torch::Tensor const& lyr_nrm_gamma_weights,
+                               torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                                  );
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& lyr_nrm_results,
+                               torch::Tensor const& lyr_nrm_mean,
+                               torch::Tensor const& lyr_nrm_invvar,
+                               torch::Tensor const& inputs, 
+                               torch::Tensor const& lyr_nrm_gamma_weights,
+                               torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               torch::Tensor const& dropout_add_mask,
+                               float                dropout_prob
+                                                  );
+
+// C++ interface
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> fwd(
+                               bool                 use_mask,
+                               bool                 use_time_mask,
+                               bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, 
+							   torch::Tensor const& lyr_nrm_gamma_weights,
+							   torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& pad_mask,
+                               float                dropout_prob
+                                                 )
+{
+  AT_ASSERTM(inputs.dim()                 == 3, "expected 3D tensor");
+  AT_ASSERTM(lyr_nrm_gamma_weights.dim()  == 1, "expected 1D tensor");
+  AT_ASSERTM(lyr_nrm_beta_weights.dim()   == 1, "expected 1D tensor");
+  AT_ASSERTM(input_weights.dim()          == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim()         == 2, "expected 2D tensor");
+
+  AT_ASSERTM(inputs.type().scalarType()                == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_gamma_weights.type().scalarType() == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_beta_weights.type().scalarType()  == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(input_weights.type().scalarType()         == at::ScalarType::Half, "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType()        == at::ScalarType::Half, "Only HALF is supported");
+
+  if (use_mask) {
+  	AT_ASSERTM(pad_mask.dim()                     == 2,                    "expected 2D tensor");
+  	AT_ASSERTM(pad_mask.type().scalarType()       == at::ScalarType::Byte, "Only BYTE is supported");
+  }
+  
+  return fwd_cuda(
+                                 use_time_mask,
+                                 is_training,
+                                 heads, 
+                                 inputs,
+                                 lyr_nrm_gamma_weights,
+                                 lyr_nrm_beta_weights,
+                                 input_weights, 
+                                 output_weights, 
+                                 use_mask ? static_cast<const uint8_t*>(pad_mask.data_ptr()) : nullptr, 
+                                 dropout_prob
+                                );
+}
+
+
+std::vector<torch::Tensor> bwd(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& lyr_nrm_results,
+                               torch::Tensor const& lyr_nrm_mean,
+                               torch::Tensor const& lyr_nrm_invvar,
+                               torch::Tensor const& inputs, 
+							   torch::Tensor const& lyr_nrm_gamma_weights,
+							   torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               torch::Tensor const& dropout_add_mask,
+                               float                dropout_prob
+                                                  )
+{
+  AT_ASSERTM(output_grads.dim()          == 3, "expected 3D tensor");
+  AT_ASSERTM(matmul2_results.dim()       == 3, "expected 3D tensor");
+  AT_ASSERTM(dropout_results.dim()       == 3, "expected 3D tensor");
+  AT_ASSERTM(softmax_results.dim()       == 3, "expected 3D tensor");
+  AT_ASSERTM(input_lin_results.dim()     == 3, "expected 3D tensor");
+  AT_ASSERTM(lyr_nrm_results.dim()       == 3, "expected 3D tensor");
+  AT_ASSERTM(lyr_nrm_mean.dim()          == 1, "expected 1D tensor");
+  AT_ASSERTM(lyr_nrm_invvar.dim()        == 1, "expected 1D tensor");
+  AT_ASSERTM(inputs.dim()                == 3, "expected 3D tensor");
+  AT_ASSERTM(lyr_nrm_gamma_weights.dim() == 1, "expected 1D tensor");
+  AT_ASSERTM(lyr_nrm_beta_weights.dim()  == 1, "expected 1D tensor");
+  AT_ASSERTM(input_weights.dim()         == 2, "expected 2D tensor");
+  AT_ASSERTM(output_weights.dim()        == 2, "expected 2D tensor");
+  AT_ASSERTM(dropout_mask.dim()          == 3, "expected 3D tensor");
+  AT_ASSERTM(dropout_add_mask.dim()      == 3, "expected 3D tensor");
+  
+  AT_ASSERTM(output_grads.type().scalarType()          == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(matmul2_results.type().scalarType()       == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(dropout_results.type().scalarType()       == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(softmax_results.type().scalarType()       == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(input_lin_results.type().scalarType()     == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_results.type().scalarType()       == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_mean.type().scalarType()          == at::ScalarType::Float, "Only FLOAT is supported");
+  AT_ASSERTM(lyr_nrm_invvar.type().scalarType()        == at::ScalarType::Float, "Only FLOAT is supported");
+  AT_ASSERTM(inputs.type().scalarType()                == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_gamma_weights.type().scalarType() == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(lyr_nrm_beta_weights.type().scalarType()  == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(input_weights.type().scalarType()         == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(output_weights.type().scalarType()        == at::ScalarType::Half,  "Only HALF is supported");
+  AT_ASSERTM(dropout_mask.type().scalarType()          == at::ScalarType::Byte,  "Only BYTE is supported");
+  AT_ASSERTM(dropout_add_mask.type().scalarType()      == at::ScalarType::Byte,  "Only BYTE is supported");
+  
+  return bwd_cuda(heads, 
+                                 output_grads,
+                                 matmul2_results,
+                                 dropout_results,
+                                 softmax_results, 
+                                 input_lin_results, 
+                                 lyr_nrm_results,
+                                 lyr_nrm_mean,
+                                 lyr_nrm_invvar,
+                                 inputs, 
+							     lyr_nrm_gamma_weights,
+								 lyr_nrm_beta_weights,
+                                 input_weights,
+                                 output_weights,
+                                 dropout_mask, 
+                                 dropout_add_mask,
+                                 dropout_prob
+                                );
+}
+
+} // end namespace cublas_gemmex
+} // end namespace self_norm_add 
+} // end namespace multihead_attn
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &multihead_attn::self_norm_add::cublas_gemmex::fwd, "Self Multihead Attention Plus Layer Norm and Residual Add Forward.");
+  m.def("backward", &multihead_attn::self_norm_add::cublas_gemmex::bwd, "Self Multihead Attention Plus Layer Norm and Residual Add Backward.");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add_cuda.cu
new file mode 100644
index 0000000000..4bf36f4edb
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add_cuda.cu
@@ -0,0 +1,556 @@
+#include <vector>
+#include <math.h>
+#include <iostream>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+
+#include "strided_batched_gemm.h"
+#include "softmax.h"
+#include "dropout.h"
+#include "layer_norm.h"
+
+// symbol to be automatically resolved by PyTorch libs
+extern THCState *state;
+
+namespace multihead_attn {
+namespace self_norm_add {
+namespace cublas_gemmex {
+
+std::vector<torch::Tensor> fwd_cuda(
+                               bool                 use_time_mask,
+							   bool                 is_training,
+                               int                  heads,
+                               torch::Tensor const& inputs, 
+							   torch::Tensor const& lyr_nrm_gamma_weights,
+							   torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               const uint8_t*       pad_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim      = inputs.size(2);
+  const int   sequences      = inputs.size(1);
+  const int   q_seq_len      = inputs.size(0);
+  const int   k_seq_len      = q_seq_len;
+  const int   batches        = sequences * q_seq_len;
+  const int   total_tokens   = batches * embed_dim;
+  const int   head_dim       = embed_dim / heads;
+  const int   output_lin_dim = 3 * embed_dim;
+  const int   attn_batches   = heads * sequences;
+  const int   lead_dim       = attn_batches * 3 * head_dim;
+  const int   batch_stride   = 3 * head_dim;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+  const float alpha          = 1.0;
+  const float beta           = 0.0;
+  const float scale          = 1.0 / sqrt(static_cast<float>(head_dim));
+ 
+  // There is no reason to use more than one stream as every kernel is 
+  // sequentially dependent
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+
+  // 3 Intermediate Results + Output (Note: dropout intermediates are generated by ATen library code)
+  auto act_options                = inputs.options().requires_grad(false);
+  auto lyr_nrm_options            = act_options.dtype(torch::kFloat32);
+  auto mask_options               = act_options.dtype(torch::kUInt8);
+
+  torch::Tensor lyr_nrm_mean      = torch::empty({batches}, lyr_nrm_options);
+  torch::Tensor lyr_nrm_invvar    = torch::empty({batches}, lyr_nrm_options);
+  torch::Tensor lyr_nrm_results   = torch::empty_like(inputs, act_options);
+
+  torch::Tensor input_lin_results = torch::empty({q_seq_len, sequences, output_lin_dim}, act_options);
+  torch::Tensor softmax_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_results   = torch::empty({attn_batches, q_seq_len, k_seq_len},   act_options);
+  torch::Tensor dropout_mask      = torch::empty({attn_batches, q_seq_len, k_seq_len},   mask_options);
+  torch::Tensor matmul2_results   = torch::empty({q_seq_len, attn_batches, head_dim},    act_options);
+  torch::Tensor output_lin_results= torch::empty_like(inputs, act_options);
+  torch::Tensor dropout_add_mask  = torch::empty_like(inputs, mask_options);
+  torch::Tensor outputs           = torch::empty_like(inputs, act_options);
+
+  // Input Linear Results Pointers to Q, K, and V of interviewed activations
+  void* q_lin_results_ptr   = static_cast<void*>(input_lin_results.data_ptr());
+  void* k_lin_results_ptr   = static_cast<void*>(static_cast<half*>(input_lin_results.data_ptr()) + head_dim);
+  void* v_lin_results_ptr   = static_cast<void*>(static_cast<half*>(input_lin_results.data_ptr()) + 2*head_dim);
+
+  // Softmax Intermediate Result Ptr (used by Matmul1 -> Softmax)
+  void* softmax_results_ptr = static_cast<void*>(softmax_results.data_ptr());
+  
+  char a_layout_t{'t'};
+  char a_layout_n{'n'};
+  char b_layout_n{'n'};
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+  // Layer Norm
+  HostApplyLayerNorm<at::Half,float>(
+                             static_cast<at::Half*>(lyr_nrm_results.data_ptr()),
+                             static_cast<float*>(lyr_nrm_mean.data_ptr()),
+                             static_cast<float*>(lyr_nrm_invvar.data_ptr()),
+                             static_cast<const at::Half*>(inputs.data_ptr()),
+                             static_cast<int>(batches),   // n1
+                             static_cast<int>(embed_dim), // n2
+                             1.0e-5,
+							 static_cast<const at::Half*>(lyr_nrm_gamma_weights.data_ptr()),
+							 static_cast<const at::Half*>(lyr_nrm_beta_weights.data_ptr()));
+
+  // Input Linear Fwd
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             output_lin_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             //static_cast<const void*>(inputs.data_ptr()),
+                             static_cast<const void*>(lyr_nrm_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             q_lin_results_ptr,
+                             CUDA_R_16F, 
+                             output_lin_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  // MatMul1 of Dot-Product Attention Plus scaling by 1/Sqrt(head size)
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             scale, 
+                             static_cast<const half*>(k_lin_results_ptr), 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<const half*>(q_lin_results_ptr),
+                             lead_dim, 
+                             batch_stride, 
+                             beta, 
+                             static_cast<half*>(softmax_results_ptr), 
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             attn_batches);
+
+  // Padded Softmax
+  bool softmax_success = false;
+  if (pad_mask == nullptr) {
+    softmax_success = dispatch_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  } else {
+    if (use_time_mask) {
+      softmax_success = dispatch_time_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             q_seq_len);
+    } else {
+      softmax_success = dispatch_masked_softmax<half, half, float>(
+                             reinterpret_cast<half*>(softmax_results_ptr),
+                             reinterpret_cast<const half*>(softmax_results_ptr),
+                             pad_mask,
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len,
+                             attn_batches*q_seq_len/sequences);
+    }
+  }
+  assert(softmax_success);
+
+  if (is_training) {
+    apex_fused_dropout_cuda<at::Half,float,uint32_t>(
+                             static_cast<at::Half const*>(softmax_results.data_ptr()), 
+                             static_cast<at::Half*>(dropout_results.data_ptr()), 
+                             static_cast<uint8_t*>(dropout_mask.data_ptr()),
+                             dropout_elems,
+                             (1.0f - dropout_prob));
+  }
+
+  // Matmul2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr), 
+                             lead_dim, 
+                             batch_stride, 
+                             (is_training) ? static_cast<const half*>(dropout_results.data_ptr()) : static_cast<const half*>(softmax_results.data_ptr()) , 
+                             //static_cast<const half*>(dropout_results.data_ptr()), 
+                             k_seq_len,  
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             static_cast<half*>(matmul2_results.data_ptr()),  
+                             head_dim*attn_batches,  
+                             head_dim, 
+                             attn_batches);
+
+  // Output Linear
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_T, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_lin_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // End-of-block Dropout-Add 
+  if (is_training) {
+    apex_dropout_add_cuda<at::Half,float,uint32_t>(
+                               static_cast<at::Half const*>(output_lin_results.data_ptr()), 
+                               static_cast<at::Half const*>(inputs.data_ptr()), 
+                               static_cast<at::Half*>(outputs.data_ptr()), 
+                               static_cast<uint8_t*>(dropout_add_mask.data_ptr()),
+                               total_tokens,
+                               (1.0f - dropout_prob));
+  } else {
+    apex_add_cuda<at::Half,float,uint32_t>(
+                               static_cast<at::Half const*>(output_lin_results.data_ptr()), 
+                               static_cast<at::Half const*>(inputs.data_ptr()), 
+                               static_cast<at::Half*>(outputs.data_ptr()), 
+                               total_tokens);
+  }
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return { 
+           lyr_nrm_results,
+		   lyr_nrm_mean,
+           lyr_nrm_invvar, 
+           input_lin_results, 
+           softmax_results,
+           dropout_results, 
+           dropout_mask, 
+           matmul2_results,
+           dropout_add_mask, 
+           outputs
+         };
+}
+
+std::vector<torch::Tensor> bwd_cuda(
+                               int                  heads,
+                               torch::Tensor const& output_grads, 
+                               torch::Tensor const& matmul2_results,
+                               torch::Tensor const& dropout_results,
+                               torch::Tensor const& softmax_results,
+                               torch::Tensor const& input_lin_results,
+                               torch::Tensor const& lyr_nrm_results,
+                               torch::Tensor const& lyr_nrm_mean,
+                               torch::Tensor const& lyr_nrm_invvar,
+                               torch::Tensor const& inputs, 
+							   torch::Tensor const& lyr_nrm_gamma_weights,
+							   torch::Tensor const& lyr_nrm_beta_weights,
+                               torch::Tensor const& input_weights,
+                               torch::Tensor const& output_weights,
+                               torch::Tensor const& dropout_mask,
+                               torch::Tensor const& dropout_add_mask,
+                               float                dropout_prob
+                                   ) 
+{
+  const int   embed_dim      = inputs.size(2);
+  const int   sequences      = inputs.size(1);
+  const int   q_seq_len      = inputs.size(0);
+  const int   k_seq_len      = q_seq_len;
+  const int   batches        = sequences * q_seq_len;
+  const int   total_tokens   = batches * embed_dim;
+  const int   head_dim       = embed_dim / heads;
+  const int   output_lin_dim = 3 * embed_dim;
+  const int   attn_batches   = heads * sequences;
+  const int   lead_dim       = attn_batches * 3 * head_dim;
+  const int   batch_stride   = 3 * head_dim;
+  const int   dropout_elems  = attn_batches * q_seq_len * k_seq_len;
+  const float alpha          = 1.0;
+  const float beta           = 0.0;
+  const float scale          = 1.0 / sqrt(static_cast<float>(head_dim));
+
+  // TODO: Streams can be used in Backprop but I haven't added more than one
+  // in my first attempt to create the code
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+  cublasSetStream(handle, stream);
+ 
+  // Output Tensor Allocations
+  torch::Tensor input_grads            = torch::empty_like(inputs);
+  torch::Tensor lyr_nrm_gamma_grads    = torch::empty_like(lyr_nrm_gamma_weights);
+  torch::Tensor lyr_nrm_beta_grads     = torch::empty_like(lyr_nrm_beta_weights);
+  torch::Tensor input_weight_grads     = torch::empty_like(input_weights);
+  torch::Tensor output_weight_grads    = torch::empty_like(output_weights);
+  // Intermediate Tensor Allocations
+  torch::Tensor dropout_add_grads      = torch::empty_like(output_grads);
+  torch::Tensor output_lin_grads       = torch::empty_like(matmul2_results);
+  torch::Tensor matmul2_grads          = torch::empty_like(dropout_results);
+  torch::Tensor input_lin_output_grads = torch::empty_like(input_lin_results);
+  torch::Tensor input_lin_grads        = torch::empty_like(inputs);
+ 
+  auto q_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr());
+  auto k_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr()) + head_dim;
+  auto v_lin_results_ptr = static_cast<half*>(input_lin_results.data_ptr()) + 2*head_dim;
+  
+  auto q_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr());
+  auto k_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr()) + head_dim;
+  auto v_lin_grads_ptr = static_cast<half*>(input_lin_output_grads.data_ptr()) + 2*head_dim;
+
+  char a_layout_n{'n'};
+  char a_layout_t{'t'};
+  char b_layout_n{'n'};
+  char b_layout_t{'t'}; 
+  
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+
+  // Dropout Add Backward  
+  apex_masked_scale_cuda<at::Half,float,uint32_t>(
+                             static_cast<at::Half const*>(output_grads.data_ptr()),
+							 static_cast<at::Half*>(dropout_add_grads.data_ptr()),
+							 static_cast<uint8_t const*>(dropout_add_mask.data_ptr()),
+       						 total_tokens,
+                             (1.0 / (1.0 - dropout_prob)));
+ 
+  // Output Linear Dgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim, 
+                             batches, 
+                             embed_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(output_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(dropout_add_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_lin_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+ 
+  // Output Linear Wgrad
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             embed_dim,
+                             batches, 
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(matmul2_results.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(dropout_add_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim, 
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(output_weight_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // MatMul2 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_t, 
+                             b_layout_n, 
+                             k_seq_len,
+                             q_seq_len,
+                             head_dim,
+                             alpha, 
+                             static_cast<const half*>(v_lin_results_ptr),
+                             lead_dim, 
+                             batch_stride,
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             beta, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len,
+                             attn_batches);
+  
+  // Matmul2 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             alpha, 
+                             static_cast<const half*>(output_lin_grads.data_ptr()),
+                             head_dim*attn_batches, 
+                             head_dim, 
+                             static_cast<const half*>(dropout_results.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             v_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+
+  // Apply Dropout Mask and Scale by Dropout Probability 
+  apex_masked_scale_cuda<at::Half,float,uint32_t>(
+                             static_cast<at::Half const*>(matmul2_grads.data_ptr()),
+							 static_cast<at::Half*>(matmul2_grads.data_ptr()),
+							 static_cast<uint8_t const*>(dropout_mask.data_ptr()),
+							 dropout_elems,
+                             (1.0 / (1.0 - dropout_prob)));
+
+  // Softmax Grad
+  bool softmax_success = false;
+  softmax_success = dispatch_softmax_backward<half, half, float>(
+                             static_cast<half*>(matmul2_grads.data_ptr()), 
+                             static_cast<half*>(matmul2_grads.data_ptr()), 
+                             reinterpret_cast<half const*>(softmax_results.data_ptr()),
+                             k_seq_len,
+                             k_seq_len,
+                             attn_batches*q_seq_len);
+  assert(softmax_success);
+
+  // Matmul1 Dgrad1
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_n, 
+                             head_dim, 
+                             q_seq_len, 
+                             k_seq_len, 
+                             scale, 
+                             k_lin_results_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             q_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+  
+  // Matmul1 Dgrad2
+  gemm_switch_fp32accum(     state, 
+                             a_layout_n, 
+                             b_layout_t, 
+                             head_dim, 
+                             k_seq_len, 
+                             q_seq_len, 
+                             scale, 
+                             q_lin_results_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             static_cast<half*>(matmul2_grads.data_ptr()),
+                             k_seq_len, 
+                             k_seq_len*q_seq_len, 
+                             beta, 
+                             k_lin_grads_ptr, 
+                             lead_dim, 
+                             batch_stride, 
+                             attn_batches);
+
+  // Input Linear Dgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_N,
+                             embed_dim,
+                             batches, 
+                             output_lin_dim,
+                             static_cast<const void*>(&alpha),
+                             static_cast<const void*>(input_weights.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F, 
+                             output_lin_dim, 
+                             static_cast<const void*>(&beta),
+                             //static_cast<void*>(input_grads.data_ptr()),
+                             static_cast<void*>(input_lin_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             //CUBLAS_GEMM_ALGO10_TENSOR_OP));
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+  
+  // Input Linear Wgrad  
+  TORCH_CUDABLAS_CHECK(cublasGemmEx(handle,
+                             CUBLAS_OP_N, 
+                             CUBLAS_OP_T,
+                             embed_dim, 
+                             output_lin_dim,
+                             batches, 
+                             static_cast<const void*>(&alpha),
+                             //static_cast<const void*>(inputs.data_ptr()),
+                             static_cast<const void*>(lyr_nrm_results.data_ptr()),
+                             CUDA_R_16F,
+                             embed_dim,
+                             static_cast<const void*>(q_lin_grads_ptr),
+                             CUDA_R_16F,
+                             output_lin_dim,
+                             static_cast<const void*>(&beta),
+                             static_cast<void*>(input_weight_grads.data_ptr()),
+                             CUDA_R_16F, 
+                             embed_dim,
+                             CUDA_R_32F,
+                             CUBLAS_GEMM_DEFAULT_TENSOR_OP));
+
+  // Fused Layer Norm Bwd with Residual Add
+  HostLayerNormGradient<half,float>(
+                             static_cast<const half*>(input_lin_grads.data_ptr()),
+                             static_cast<half const*>(output_grads.data_ptr()), 
+                             static_cast<const float*>(lyr_nrm_mean.data_ptr()),
+                             static_cast<const float*>(lyr_nrm_invvar.data_ptr()),
+                             inputs,
+                             static_cast<int>(batches),   // n1
+                             static_cast<int>(embed_dim), // n2
+                             static_cast<const half*>(lyr_nrm_gamma_weights.data_ptr()),
+                             static_cast<const half*>(lyr_nrm_beta_weights.data_ptr()),
+                             1.0e-5,
+                             static_cast<half*>(input_grads.data_ptr()),
+                             static_cast<half*>(lyr_nrm_gamma_grads.data_ptr()),
+                             static_cast<half*>(lyr_nrm_beta_grads.data_ptr())
+                                   );
+
+  TORCH_CUDABLAS_CHECK(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+
+  return {
+		   input_grads, 
+           lyr_nrm_gamma_grads, 
+           lyr_nrm_beta_grads, 
+           input_weight_grads, 
+           output_weight_grads
+         };
+}
+
+} // end namespace cublas_gemmex
+} // end namespace self_norm_add 
+} // end namespace multihead_attn
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/softmax.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/softmax.h
new file mode 100644
index 0000000000..e669fd4e6a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/softmax.h
@@ -0,0 +1,2641 @@
+#pragma once
+#include <ATen/CUDAGeneratorImpl.h>
+#include <ATen/cuda/CUDAGraphsUtils.cuh>
+#include <curand_kernel.h>
+#include "philox.h"
+ 
+#include <assert.h>
+#include <cfloat>
+#include <limits>
+#include <stdint.h>
+#include <cuda_fp16.h>
+#include <cmath>
+ 
+namespace {
+    template <typename Datatype, int ELEMENTS_PER_LDG>
+    __device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);
+ 
+    template <>
+    __device__ __inline__ void copy_vector<__half, 1>(__half *dst, const __half *src) { *dst = *src; }
+ 
+    template <>
+    __device__ __inline__ void copy_vector<float, 1>(float *dst, const float *src) { *dst = *src; }
+ 
+    template <>
+    __device__ __inline__ void copy_vector<__half, 4>(__half *dst, const __half *src) { *((float2*) dst) = *((float2*) src); } 
+    template <>
+    __device__ __inline__ void copy_vector<uint8_t, 1>(uint8_t *dst, const uint8_t *src) { *dst = *src; }
+    
+    template <>
+    __device__ __inline__ void copy_vector<uint8_t, 4>(uint8_t *dst, const uint8_t *src) {*((half2*) dst) = *((half2*) src); }
+   
+    template <typename Datatype, int ELEMENTS_PER_LDG>
+    __device__ __inline__ void apply_mask(Datatype *dst, Datatype value, const uint8_t *src);
+    
+    template <>
+    __device__ __inline__ void apply_mask<__half, 1>(__half *dst, __half value, const uint8_t *src) {
+      if (*src == 1) { *dst = value; }
+    }
+    template <typename Datatype, int ELEMENTS_PER_LDG>
+    __device__ __inline__ void apply_additive_mask(Datatype *dst, const Datatype *additive_mask);
+    template <>
+    __device__ __inline__ void apply_additive_mask<__half, 1>(__half *dst, const __half *additive_mask) {
+      *dst += *additive_mask; 
+    }
+    template <>
+    __device__ __inline__ void apply_additive_mask<__half, 4>(__half *dst, const __half *additive_mask) {
+      *dst += *additive_mask;
+      *(dst+1) += *(additive_mask+1);
+      *(dst+2) += *(additive_mask+2);
+      *(dst+3) += *(additive_mask+3);}    
+} // namespace anonymous
+
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+// Warp Softmax forward
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t, typename acc_t, int WARP_BATCH, int WARP_ITERATIONS, int WARP_SIZE = 32, int ELEMENTS_PER_LDG_STG=1>
+__global__ void softmax_warp_forward(input_t *dst, const output_t *src, int batch_size, int stride, int element_count)
+{
+    assert(ELEMENTS_PER_LDG_STG==1);
+ 
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+ 
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+ 
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+ 
+    src += first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+    dst += first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+ 
+    // load data from global memory
+    input_t elements_input[WARP_BATCH][WARP_ITERATIONS];
+    for (int i = 0;i < WARP_BATCH;++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            #pragma unroll
+            for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                elements_input[i][it + element] = -std::numeric_limits<float>::infinity();
+            }
+ 
+            if (element_index < batch_element_count) {
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], src + i * element_count + it * WARP_SIZE);
+            }
+ 
+        }
+    }
+ 
+    // convert input_t to acc_t
+    acc_t elements[WARP_BATCH][WARP_ITERATIONS];
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            elements[i][it] = elements_input[i][it];
+        }
+    }
+ 
+    constexpr uint32_t  FULL_MASK = 0xffffffff;
+ 
+    // compute local max_value
+ 
+    // take the max_value of the first element to avoid one max call
+    acc_t max_value[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        max_value[i] = elements[i][0];
+    }
+ 
+    #pragma unroll
+    for (int it = 1;it < WARP_ITERATIONS;++it) {
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
+        }
+    }
+ 
+    // reduction max_value
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        float val[WARP_BATCH];
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);
+        }
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];
+        }
+    }
+ 
+    // compute local sum
+    acc_t sum[WARP_BATCH] { 0.0f };
+ 
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            //elements[i][it] = expf(elements[i][it] - max_value[i]);
+            elements[i][it] = std::exp(elements[i][it] - max_value[i]);
+            sum[i] += elements[i][it];
+        }
+    }
+ 
+    // reduction sum
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);
+        }
+    }
+ 
+    // store result
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                //dst[i * element_count + it * WARP_SIZE] = elements[i][it] / sum[i];
+                output_t out[ELEMENTS_PER_LDG_STG];
+                for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                    out[element] = elements[i][it + element] / sum[i];
+                }
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(dst + i * element_count + it * WARP_SIZE, out);
+            }
+            else {
+                break;
+            }
+        }
+    }
+}
+
+
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t>
+using softmax_forward_func = void(*)(input_t *dst, const output_t *src, int batch_size, int stride, int element_count);
+ 
+template <typename input_t, typename output_t, typename acc_t>
+bool warp_softmax_kernel(int log2_elements, int &warp_size, int &batches_per_warp, softmax_forward_func<input_t, output_t> &kernel) {
+    // determine size of a warp
+    const int next_power_of_two = 1 << log2_elements;
+    warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;
+ 
+    // determine how many batches a warp should process.
+    batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+ 
+    switch (log2_elements) {
+    case 0: // 1
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 2,1,1,1>;
+        break;
+    case 1: // 2
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 2,1,2,1>;
+        break;
+    case 2: // 4
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 2,1,4,1>;
+        break;
+    case 3: // 8
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 2,1,8,1>;
+        break;
+    case 4: // 16
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 2,1,16,1>;
+        break;
+    case 5: // 32
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 2,1,32,1>;
+        break;
+    case 6: // 64
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 2,2,32,1>;
+        break;
+    case 7: // 128
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 2,4,32,1>;
+        break;
+    case 8: // 256
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 1,8,32,1>;
+        break;
+    case 9: // 512
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 1,16,32,1>;
+        break;
+    case 10: // 1024
+        kernel = &softmax_warp_forward<input_t, output_t, acc_t, 1,32,32,1>;
+        break;
+    default:
+        return false;
+    }
+    return true;
+}
+
+template<typename input_t, typename output_t, typename acc_t>
+bool dispatch_softmax(output_t *dst, const input_t *src, int softmax_elements, int softmax_elements_stride, int batch_count)
+{
+    if (softmax_elements == 0) {
+        return true;
+    } else if (softmax_elements <= 1024) {
+        // compute function index. there's a function for each power of two size up to 1024.
+        int log2_elements = 0;
+        while ((1 << log2_elements) < softmax_elements) ++log2_elements;
+ 
+        softmax_forward_func<input_t, output_t> kernel;
+        int warp_size, batches_per_warp;
+        if (!warp_softmax_kernel<input_t, output_t, acc_t>(log2_elements, warp_size, batches_per_warp, kernel)) {
+            return false;
+        }
+ 
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+ 
+        // compute warps per block.
+        int warps_per_block = (threads_per_block / warp_size);
+ 
+        // compute launch size
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+ 
+        // launch
+        kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, batch_count, softmax_elements_stride, softmax_elements);
+        return true;
+    }
+    return false;
+}
+
+template <typename input_t, typename output_t, typename acc_t, int WARP_BATCH, int WARP_ITERATIONS, int WARP_SIZE, int ELEMENTS_PER_LDG_STG>
+__global__ void additive_masked_softmax_dropout_warp_forward_vec4(output_t *dst, uint8_t *dropout_mask, const input_t *src, const input_t *pad_mask, int batch_size, int stride, int element_count, int pad_batch_stride, at::PhiloxCudaState philox_args, float p)
+{
+ 
+    assert(ELEMENTS_PER_LDG_STG==4);
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+    int tid = blockIdx.x * blockDim.x * blockDim.y + threadIdx.y * blockDim.x + threadIdx.x;
+    acc_t pinv = acc_t(1)/p;
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+     
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+    //vectorize if element_count is multiple of 4, else don't vectorize
+    input_t elements_input[WARP_BATCH][WARP_ITERATIONS];
+
+    int thread_offset =  first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+    src += thread_offset;
+    dst += thread_offset;
+    dropout_mask += thread_offset;
+    
+    // load data from global memory
+    for (int i = 0;i < WARP_BATCH;++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        int pad_thread_offset = ( (first_batch + i) / pad_batch_stride) * stride + ELEMENTS_PER_LDG_STG * local_idx;
+        const half* curr_mask    = pad_mask + pad_thread_offset;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            #pragma unroll
+            for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+    	//masking_value is a large negative value
+                elements_input[i][it + element] = -10000;
+            }
+    
+            if (element_index < batch_element_count) {
+                int itr_jmp = it * WARP_SIZE;
+                int itr_idx = i * element_count + itr_jmp;
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], src + itr_idx);
+                apply_additive_mask<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], curr_mask + itr_jmp); //(__half)-std::numeric_limits<float>::infinity()
+            } 
+    
+        }
+    }
+    // convert input_t to acc_t
+    acc_t elements[WARP_BATCH][WARP_ITERATIONS];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            elements[i][it] = elements_input[i][it];
+        }
+    }
+ 
+    constexpr uint32_t  FULL_MASK = 0xffffffff;
+ 
+    // compute local max_value
+ 
+    // take the max_value of the first element to avoid one max call
+    acc_t max_value[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        max_value[i] = elements[i][0];
+    }
+ 
+    #pragma unroll
+    for (int it = 1;it < WARP_ITERATIONS;++it) {
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
+        }
+    }
+ 
+    // reduction max_value
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        float val[WARP_BATCH];
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);
+        }
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];
+        }
+    }
+ 
+    // compute local sum
+    acc_t sum[WARP_BATCH] { 0.0f };
+ 
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            elements[i][it] = std::exp(elements[i][it] - max_value[i]);
+            sum[i] += elements[i][it];
+        }
+    }
+ 
+    // reduction sum
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);
+        }
+    }
+    auto seeds = at::cuda::philox::unpack(philox_args);
+    Philox ph(std::get<0>(seeds), tid, std::get<1>(seeds));     
+    uint8_t rands[WARP_BATCH][WARP_ITERATIONS];
+    float4 rand_num;
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        if (i >= local_batches)
+            break;
+	#pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it+=ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+		rand_num = uniform4(ph());
+                rands[i][it] = (rand_num.x <= p) > 0.5;  
+                rands[i][it+1] = (rand_num.y <= p) > 0.5;
+                rands[i][it+2] = (rand_num.z <= p) > 0.5;
+                rands[i][it+3] = (rand_num.w <= p) > 0.5;
+                copy_vector<uint8_t, ELEMENTS_PER_LDG_STG>(dropout_mask + i * element_count + it * WARP_SIZE, &rands[i][it]);
+	    }
+        }
+    }
+
+    // store result
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                output_t out[ELEMENTS_PER_LDG_STG];
+                #pragma unroll
+                for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                    out[element] = rands[i][it+element] * (pinv * (elements[i][it + element] / sum[i]));
+                }
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(dst + i * element_count + it * WARP_SIZE, out);
+    
+            }
+            else {
+                break;
+            }
+        }
+    }
+}
+
+
+template <typename input_t, typename output_t, typename acc_t, int WARP_BATCH, int WARP_ITERATIONS, int WARP_SIZE, int ELEMENTS_PER_LDG_STG>
+__global__ void additive_masked_softmax_dropout_warp_forward(output_t *dst, uint8_t *dropout_mask, const input_t *src, const input_t *pad_mask, int batch_size, int stride, int element_count, int pad_batch_stride, at::PhiloxCudaState philox_args, float p)
+{
+    assert(ELEMENTS_PER_LDG_STG==1);
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+    int tid = blockIdx.x * blockDim.x * blockDim.y + threadIdx.y * blockDim.x + threadIdx.x;
+    acc_t pinv = acc_t(1)/p;
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+     
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+    //vectorize if element_count is multiple of 4, else don't vectorize
+    input_t elements_input[WARP_BATCH][WARP_ITERATIONS];
+
+    int thread_offset =  first_batch * stride + local_idx;
+    src += thread_offset;
+    dst += thread_offset;
+    dropout_mask += thread_offset;
+    
+    // load data from global memory
+    for (int i = 0;i < WARP_BATCH;++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        int pad_thread_offset = ( (first_batch + i) / pad_batch_stride) * stride + local_idx;
+        const half* curr_mask    = pad_mask + pad_thread_offset;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += 1) {
+            int element_index = local_idx + it * WARP_SIZE;
+            #pragma unroll
+            for (int element = 0;element < 1;++element) {
+    	//masking_value is a large negative value
+                elements_input[i][it + element] = -10000;
+            }
+    
+            if (element_index < batch_element_count) {
+                int itr_jmp = it * WARP_SIZE;
+                int itr_idx = i * element_count + itr_jmp;
+                copy_vector<input_t, 1>(&elements_input[i][it], src + itr_idx);
+                apply_additive_mask<input_t, 1>(&elements_input[i][it], curr_mask + itr_jmp); 
+            } 
+    
+        }
+    }
+    // convert input_t to acc_t
+    acc_t elements[WARP_BATCH][WARP_ITERATIONS];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            elements[i][it] = elements_input[i][it];
+        }
+    }
+ 
+    constexpr uint32_t  FULL_MASK = 0xffffffff;
+ 
+    // compute local max_value
+ 
+    // take the max_value of the first element to avoid one max call
+    acc_t max_value[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        max_value[i] = elements[i][0];
+    }
+ 
+    #pragma unroll
+    for (int it = 1;it < WARP_ITERATIONS;++it) {
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
+        }
+    }
+ 
+    // reduction max_value
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        float val[WARP_BATCH];
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);
+        }
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];
+        }
+    }
+ 
+    // compute local sum
+    acc_t sum[WARP_BATCH] { 0.0f };
+ 
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            elements[i][it] = std::exp(elements[i][it] - max_value[i]);
+            sum[i] += elements[i][it];
+        }
+    }
+ 
+    // reduction sum
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);
+        }
+    }
+    curandStatePhilox4_32_10_t state;
+    auto seeds = at::cuda::philox::unpack(philox_args);
+    curand_init(
+      std::get<0>(seeds),
+      tid,
+      std::get<1>(seeds),
+      &state);
+     
+    // store result
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += 1) {
+            int element_index = local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                output_t out[1];
+                acc_t softmax_out[1];
+                uint8_t dropout_mask_temp[1];
+                //generate a vector of random numbers here 
+                float rand = curand_uniform(&state);
+                float *rand_ptr = (float*)(&rand);    
+                #pragma unroll
+                for (int element = 0;element < 1;++element) {
+    	        softmax_out[element] = (elements[i][it + element] / sum[i]);	
+                    rand_ptr[element] = rand_ptr[element] <= p;       
+                    out[element] = rand_ptr[element] * pinv * softmax_out[element];
+    	            dropout_mask_temp[element] = rand_ptr[element] > 0.5; // just to distinguish 0.0f and 1.0f 
+                }
+                copy_vector<output_t, 1>(dst + i * element_count + it * WARP_SIZE, out);
+                copy_vector<uint8_t, 1>(dropout_mask + i * element_count + it * WARP_SIZE, dropout_mask_temp);
+    
+            }
+            else {
+                break;
+            }
+        }
+    }
+}
+
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t, typename acc_t>
+using additive_masked_softmax_dropout_forward_func = void(*)(output_t *dst, uint8_t *dropout_mask, const input_t *src, const input_t *pad_mask, int batch_size, int stride, int element_count, int pad_batch_stride,  at::PhiloxCudaState philox_args, float p);
+
+
+template <typename input_t, typename output_t, typename acc_t>
+bool warp_additive_masked_softmax_dropout_kernel(int element_count, int log2_elements, int &warp_size, int &batches_per_warp, additive_masked_softmax_dropout_forward_func<input_t, output_t, acc_t> &kernel) {
+    // determine size of a warp
+    const int next_power_of_two = 1 << log2_elements;
+    warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;
+ 
+    // determine how many batches a warp should process.
+    batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+    bool flag_vec4 = (element_count % 4 == 0); 
+    switch (log2_elements) {
+    case 0: // 1
+        kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 2,1,1,1>;
+        break;
+    case 1: // 2
+        kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 2,1,2,1>;
+        break;
+    case 2: // 4
+        kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 2,1,4,1>;
+        break;
+    case 3: // 8
+        kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 2,1,8,1>;
+        break;
+    case 4: // 16
+        kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 2,1,16,1>;
+        break;
+    case 5: // 32
+        kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 2,1,32,1>;
+        break;
+    case 6: // 64
+        kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 2,2,32,1>;
+        break;
+    case 7: // 128
+	if (flag_vec4) kernel = &additive_masked_softmax_dropout_warp_forward_vec4<input_t, output_t, acc_t, 2,4,32,4>;
+	else kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 2,4,32,1>;
+        break;
+    case 8: // 256
+	if (flag_vec4) kernel = &additive_masked_softmax_dropout_warp_forward_vec4<input_t, output_t, acc_t, 1,8,32,4>;
+	else kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 1,8,32,1>;
+        break;
+    case 9: // 512
+        if (flag_vec4) kernel = &additive_masked_softmax_dropout_warp_forward_vec4<input_t, output_t, acc_t, 1,16,32,4>;
+	else kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 1,16,32,1>;
+        break;
+    case 10: // 1024
+        if (flag_vec4) kernel = &additive_masked_softmax_dropout_warp_forward_vec4<input_t, output_t, acc_t, 1,32,32,4>;
+	else kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 1,32,32,1>;
+        break;
+    case 11: // 2048
+        if (flag_vec4) kernel = &additive_masked_softmax_dropout_warp_forward_vec4<input_t, output_t, acc_t, 1,64,32,4>;
+	else kernel = &additive_masked_softmax_dropout_warp_forward<input_t, output_t, acc_t, 1,64,32,1>;
+        break;
+    default:
+        return false;
+    }
+    return true;
+}
+
+
+
+template<typename input_t, typename output_t, typename acc_t>
+bool dispatch_additive_masked_softmax_dropout(output_t *dst, uint8_t *dropout_mask, const input_t *src, const input_t *pad_mask, int totalElements, int softmax_elements, int softmax_elements_stride, int batch_count, int pad_batch_stride, float p, cudaStream_t streamid)// p is the probability to keep, not drop
+{
+	
+    if (softmax_elements == 0) {
+        return true;
+    } else if (softmax_elements <= 2048) {
+        // compute function index. there's a function for each power of two size up to 1024.
+        int log2_elements = 0;
+        while ((1 << log2_elements) < softmax_elements) ++log2_elements;
+ 
+        additive_masked_softmax_dropout_forward_func<input_t, output_t, acc_t> kernel;
+        int warp_size, batches_per_warp;
+        if (!warp_additive_masked_softmax_dropout_kernel<input_t, output_t, acc_t>(softmax_elements, log2_elements, warp_size, batches_per_warp, kernel)) {
+            return false;
+        }
+ 
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+        // compute warps per block.
+        int warps_per_block = (threads_per_block / warp_size);
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+	c10::optional<at::Generator> gen_;
+        auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(gen_, at::cuda::detail::getDefaultCUDAGenerator());
+        int64_t counter_offset = (totalElements/(blocks*threads_per_block)+1);
+        at::PhiloxCudaState rng_engine_inputs;
+	{
+          std::lock_guard<std::mutex> lock(gen->mutex_);
+	  rng_engine_inputs = gen->philox_cuda_state(counter_offset);
+        }
+ 
+        // compute launch size
+        dim3 threads(warp_size, warps_per_block, 1);
+         
+        // launch
+        kernel<<<blocks, threads, 0, streamid>>>(dst, dropout_mask, src, pad_mask, batch_count, softmax_elements_stride, softmax_elements, pad_batch_stride, rng_engine_inputs, p);
+        return true;
+    }
+    return false;
+}
+
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t, typename acc_t, int WARP_BATCH, int WARP_ITERATIONS, int WARP_SIZE = 32, int ELEMENTS_PER_LDG_STG=1>
+__global__ void additive_masked_softmax_warp_forward(input_t *dst, const output_t *src, const input_t *pad_mask, int batch_size, int stride, int element_count, int pad_batch_stride)
+{
+    assert(ELEMENTS_PER_LDG_STG==1);
+ 
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+ 
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+ 
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+
+    int thread_offset =  first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+    src += thread_offset;
+    dst += thread_offset;
+ 
+    // load data from global memory
+    input_t elements_input[WARP_BATCH][WARP_ITERATIONS];
+    for (int i = 0;i < WARP_BATCH;++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        int pad_thread_offset = ( (first_batch + i) / pad_batch_stride) * stride + ELEMENTS_PER_LDG_STG * local_idx;
+        const half* curr_mask    = pad_mask + pad_thread_offset;
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            #pragma unroll
+            for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+		//masking_value is a large negative value
+                elements_input[i][it + element] = -10000;
+            }
+ 
+            if (element_index < batch_element_count) {
+                int itr_jmp = it * WARP_SIZE;
+                int itr_idx = i * element_count + itr_jmp;
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], src + itr_idx);
+                //apply_mask<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], 
+                //                                          (__half)-std::numeric_limits<float>::infinity(), 
+                //                                          curr_mask + itr_jmp);
+                elements_input[i][it] += *(curr_mask + itr_jmp);
+	    } 
+ 
+        }
+    }
+ 
+    // convert input_t to acc_t
+    acc_t elements[WARP_BATCH][WARP_ITERATIONS];
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            elements[i][it] = elements_input[i][it];
+        }
+    }
+ 
+    constexpr uint32_t  FULL_MASK = 0xffffffff;
+ 
+    // compute local max_value
+ 
+    // take the max_value of the first element to avoid one max call
+    acc_t max_value[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        max_value[i] = elements[i][0];
+    }
+ 
+    #pragma unroll
+    for (int it = 1;it < WARP_ITERATIONS;++it) {
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
+        }
+    }
+ 
+    // reduction max_value
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        float val[WARP_BATCH];
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);
+        }
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];
+        }
+    }
+ 
+    // compute local sum
+    acc_t sum[WARP_BATCH] { 0.0f };
+ 
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            //elements[i][it] = expf(elements[i][it] - max_value[i]);
+            elements[i][it] = std::exp(elements[i][it] - max_value[i]);
+            sum[i] += elements[i][it];
+        }
+    }
+ 
+    // reduction sum
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);
+        }
+    }
+ 
+    // store result
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                //dst[i * element_count + it * WARP_SIZE] = elements[i][it] / sum[i];
+                output_t out[ELEMENTS_PER_LDG_STG];
+                for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                    out[element] = elements[i][it + element] / sum[i];
+                }
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(dst + i * element_count + it * WARP_SIZE, out);
+            }
+            else {
+                break;
+            }
+        }
+    }
+}
+
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t>
+using additive_masked_softmax_forward_func = void(*)(input_t *dst, const output_t *src, const half *pad_mask, int batch_size, int stride, int element_count, int pad_batch_stride);
+ 
+template <typename input_t, typename output_t, typename acc_t>
+bool warp_additive_masked_softmax_kernel(int log2_elements, int &warp_size, int &batches_per_warp, additive_masked_softmax_forward_func<input_t, output_t> &kernel) {
+    // determine size of a warp
+    const int next_power_of_two = 1 << log2_elements;
+    warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;
+ 
+    // determine how many batches a warp should process.
+    batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+ 
+    switch (log2_elements) {
+    case 0: // 1
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,1,1>;
+        break;
+    case 1: // 2
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,2,1>;
+        break;
+    case 2: // 4
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,4,1>;
+        break;
+    case 3: // 8
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,8,1>;
+        break;
+    case 4: // 16
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,16,1>;
+        break;
+    case 5: // 32
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,32,1>;
+        break;
+    case 6: // 64
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,2,32,1>;
+        break;
+    case 7: // 128
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,4,32,1>;
+        break;
+    case 8: // 256
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 1,8,32,1>;
+        break;
+    case 9: // 512
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 1,16,32,1>;
+        break;
+    case 10: // 1024
+        kernel = &additive_masked_softmax_warp_forward<input_t, output_t, acc_t, 1,32,32,1>;
+        break;
+    default:
+        return false;
+    }
+    return true;
+}
+ 
+template<typename input_t, typename output_t, typename acc_t>
+bool dispatch_additive_masked_softmax(output_t *dst, const input_t *src, const input_t *pad_mask, int softmax_elements, int softmax_elements_stride, int batch_count, int pad_batch_stride)
+{
+    if (softmax_elements == 0) {
+        return true;
+    } else if (softmax_elements <= 1024) {
+        // compute function index. there's a function for each power of two size up to 1024.
+        int log2_elements = 0;
+        while ((1 << log2_elements) < softmax_elements) ++log2_elements;
+ 
+        additive_masked_softmax_forward_func<input_t, output_t> kernel;
+        int warp_size, batches_per_warp;
+        if (!warp_additive_masked_softmax_kernel<input_t, output_t, acc_t>(log2_elements, warp_size, batches_per_warp, kernel)) {
+            return false;
+        }
+ 
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+ 
+        // compute warps per block.
+        int warps_per_block = (threads_per_block / warp_size);
+ 
+        // compute launch size
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+ 
+        // launch
+        kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, pad_mask, batch_count, softmax_elements_stride, softmax_elements, pad_batch_stride);
+        return true;
+    }
+    return false;
+}
+
+template<typename input_t, typename output_t, typename acc_t>
+bool dispatch_additive_masked_softmax_stream(output_t *dst, const input_t *src, const input_t *pad_mask, int softmax_elements, int softmax_elements_stride, int batch_count, int pad_batch_stride, cudaStream_t streamid)
+{
+    if (softmax_elements == 0) {
+        return true;
+    } else if (softmax_elements <= 1024) {
+        // compute function index. there's a function for each power of two size up to 1024.
+        int log2_elements = 0;
+        while ((1 << log2_elements) < softmax_elements) ++log2_elements;
+        additive_masked_softmax_forward_func<input_t, output_t> kernel;
+        int warp_size, batches_per_warp;
+        if (!warp_additive_masked_softmax_kernel<input_t, output_t, acc_t>(log2_elements, warp_size, batches_per_warp, kernel)) {
+            return false;
+        }
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+        // compute warps per block.
+        int warps_per_block = (threads_per_block / warp_size);
+        // compute launch size
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+        // launch
+        kernel<<<blocks, threads, 0, streamid>>>(dst, src, pad_mask, batch_count, softmax_elements_stride, softmax_elements, pad_batch_stride);
+        return true;
+    }
+    return false;
+}
+
+
+
+
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t, typename acc_t, int WARP_BATCH, int WARP_ITERATIONS, int WARP_SIZE = 32, int ELEMENTS_PER_LDG_STG=1>
+__global__ void masked_softmax_warp_forward(input_t *dst, const output_t *src, const uint8_t *pad_mask, int batch_size, int stride, int element_count, int pad_batch_stride)
+{
+    assert(ELEMENTS_PER_LDG_STG==1);
+ 
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+ 
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+ 
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+
+    int thread_offset =  first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+    src += thread_offset;
+    dst += thread_offset;
+ 
+    // load data from global memory
+    input_t elements_input[WARP_BATCH][WARP_ITERATIONS];
+    for (int i = 0;i < WARP_BATCH;++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        int pad_thread_offset = ( (first_batch + i) / pad_batch_stride) * stride + ELEMENTS_PER_LDG_STG * local_idx;
+        const uint8_t* curr_mask    = pad_mask + pad_thread_offset;
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            #pragma unroll
+            for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                elements_input[i][it + element] = -std::numeric_limits<float>::infinity();
+            }
+ 
+            if (element_index < batch_element_count) {
+                int itr_jmp = it * WARP_SIZE;
+                int itr_idx = i * element_count + itr_jmp;
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], src + itr_idx);
+                apply_mask<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], 
+                                                          (__half)-std::numeric_limits<float>::infinity(), 
+                                                          curr_mask + itr_jmp);
+            }
+ 
+        }
+    }
+ 
+    // convert input_t to acc_t
+    acc_t elements[WARP_BATCH][WARP_ITERATIONS];
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            elements[i][it] = elements_input[i][it];
+        }
+    }
+ 
+    constexpr uint32_t  FULL_MASK = 0xffffffff;
+ 
+    // compute local max_value
+ 
+    // take the max_value of the first element to avoid one max call
+    acc_t max_value[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        max_value[i] = elements[i][0];
+    }
+ 
+    #pragma unroll
+    for (int it = 1;it < WARP_ITERATIONS;++it) {
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
+        }
+    }
+ 
+    // reduction max_value
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        float val[WARP_BATCH];
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);
+        }
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];
+        }
+    }
+ 
+    // compute local sum
+    acc_t sum[WARP_BATCH] { 0.0f };
+ 
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            //elements[i][it] = expf(elements[i][it] - max_value[i]);
+            elements[i][it] = std::exp(elements[i][it] - max_value[i]);
+            sum[i] += elements[i][it];
+        }
+    }
+ 
+    // reduction sum
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);
+        }
+    }
+ 
+    // store result
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                //dst[i * element_count + it * WARP_SIZE] = elements[i][it] / sum[i];
+                output_t out[ELEMENTS_PER_LDG_STG];
+                for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                    out[element] = elements[i][it + element] / sum[i];
+                }
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(dst + i * element_count + it * WARP_SIZE, out);
+            }
+            else {
+                break;
+            }
+        }
+    }
+}
+
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t>
+using masked_softmax_forward_func = void(*)(input_t *dst, const output_t *src, const uint8_t *pad_mask, int batch_size, int stride, int element_count, int pad_batch_stride);
+ 
+template <typename input_t, typename output_t, typename acc_t>
+bool warp_masked_softmax_kernel(int log2_elements, int &warp_size, int &batches_per_warp, masked_softmax_forward_func<input_t, output_t> &kernel) {
+    // determine size of a warp
+    const int next_power_of_two = 1 << log2_elements;
+    warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;
+ 
+    // determine how many batches a warp should process.
+    batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+ 
+    switch (log2_elements) {
+    case 0: // 1
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,1,1>;
+        break;
+    case 1: // 2
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,2,1>;
+        break;
+    case 2: // 4
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,4,1>;
+        break;
+    case 3: // 8
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,8,1>;
+        break;
+    case 4: // 16
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,16,1>;
+        break;
+    case 5: // 32
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,32,1>;
+        break;
+    case 6: // 64
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 2,2,32,1>;
+        break;
+    case 7: // 128
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 2,4,32,1>;
+        break;
+    case 8: // 256
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 1,8,32,1>;
+        break;
+    case 9: // 512
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 1,16,32,1>;
+        break;
+    case 10: // 1024
+        kernel = &masked_softmax_warp_forward<input_t, output_t, acc_t, 1,32,32,1>;
+        break;
+    default:
+        return false;
+    }
+    return true;
+}
+ 
+template<typename input_t, typename output_t, typename acc_t>
+bool dispatch_masked_softmax(output_t *dst, const input_t *src, const uint8_t *pad_mask, int softmax_elements, int softmax_elements_stride, int batch_count, int pad_batch_stride)
+{
+    if (softmax_elements == 0) {
+        return true;
+    } else if (softmax_elements <= 1024) {
+        // compute function index. there's a function for each power of two size up to 1024.
+        int log2_elements = 0;
+        while ((1 << log2_elements) < softmax_elements) ++log2_elements;
+ 
+        masked_softmax_forward_func<input_t, output_t> kernel;
+        int warp_size, batches_per_warp;
+        if (!warp_masked_softmax_kernel<input_t, output_t, acc_t>(log2_elements, warp_size, batches_per_warp, kernel)) {
+            return false;
+        }
+ 
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+ 
+        // compute warps per block.
+        int warps_per_block = (threads_per_block / warp_size);
+ 
+        // compute launch size
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+ 
+        // launch
+        kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, pad_mask, batch_count, softmax_elements_stride, softmax_elements, pad_batch_stride);
+        return true;
+    }
+    return false;
+}
+
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t, typename acc_t, int WARP_BATCH, int WARP_ITERATIONS, int WARP_SIZE = 32, int ELEMENTS_PER_LDG_STG=1>
+__global__ void time_masked_softmax_warp_forward(input_t *dst, const output_t *src, const uint8_t *pad_mask, int batch_size, int stride, int element_count, int mod_seq_len)
+{
+    assert(ELEMENTS_PER_LDG_STG==1);
+ 
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+ 
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+ 
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+
+    int thread_offset =  first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+    src += thread_offset;
+    dst += thread_offset;
+ 
+    // load data from global memory
+    input_t elements_input[WARP_BATCH][WARP_ITERATIONS];
+    for (int i = 0;i < WARP_BATCH;++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        int pad_thread_offset = ( (first_batch + i) % mod_seq_len) * stride + ELEMENTS_PER_LDG_STG * local_idx;
+        const uint8_t* curr_mask    = pad_mask + pad_thread_offset;
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            #pragma unroll
+            for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                elements_input[i][it + element] = -std::numeric_limits<float>::infinity();
+            }
+ 
+            if (element_index < batch_element_count) {
+                int itr_jmp = it * WARP_SIZE;
+                int itr_idx = i * element_count + itr_jmp;
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], src + itr_idx);
+                apply_mask<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], 
+                                                          (__half)-std::numeric_limits<float>::infinity(), 
+                                                          curr_mask + itr_jmp);
+            }
+ 
+        }
+    }
+ 
+    // convert input_t to acc_t
+    acc_t elements[WARP_BATCH][WARP_ITERATIONS];
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            elements[i][it] = elements_input[i][it];
+        }
+    }
+ 
+    constexpr uint32_t  FULL_MASK = 0xffffffff;
+ 
+    // compute local max_value
+ 
+    // take the max_value of the first element to avoid one max call
+    acc_t max_value[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        max_value[i] = elements[i][0];
+    }
+ 
+    #pragma unroll
+    for (int it = 1;it < WARP_ITERATIONS;++it) {
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
+        }
+    }
+ 
+    // reduction max_value
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        float val[WARP_BATCH];
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);
+        }
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];
+        }
+    }
+ 
+    // compute local sum
+    acc_t sum[WARP_BATCH] { 0.0f };
+ 
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            //elements[i][it] = expf(elements[i][it] - max_value[i]);
+            elements[i][it] = std::exp(elements[i][it] - max_value[i]);
+            sum[i] += elements[i][it];
+        }
+    }
+ 
+    // reduction sum
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);
+        }
+    }
+ 
+    // store result
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                //dst[i * element_count + it * WARP_SIZE] = elements[i][it] / sum[i];
+                output_t out[ELEMENTS_PER_LDG_STG];
+                for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                    out[element] = elements[i][it + element] / sum[i];
+                }
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(dst + i * element_count + it * WARP_SIZE, out);
+            }
+            else {
+                break;
+            }
+        }
+    }
+}
+
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t>
+using time_masked_softmax_forward_func = void(*)(input_t *dst, const output_t *src, const uint8_t *pad_mask, int batch_size, int stride, int element_count, int mod_seq_len);
+ 
+template <typename input_t, typename output_t, typename acc_t>
+bool warp_time_masked_softmax_kernel(int log2_elements, int &warp_size, int &batches_per_warp, time_masked_softmax_forward_func<input_t, output_t> &kernel) {
+    // determine size of a warp
+    const int next_power_of_two = 1 << log2_elements;
+    warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;
+ 
+    // determine how many batches a warp should process.
+    batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+ 
+    switch (log2_elements) {
+    case 0: // 1
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,1,1>;
+        break;
+    case 1: // 2
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,2,1>;
+        break;
+    case 2: // 4
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,4,1>;
+        break;
+    case 3: // 8
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,8,1>;
+        break;
+    case 4: // 16
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,16,1>;
+        break;
+    case 5: // 32
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,1,32,1>;
+        break;
+    case 6: // 64
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,2,32,1>;
+        break;
+    case 7: // 128
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 2,4,32,1>;
+        break;
+    case 8: // 256
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 1,8,32,1>;
+        break;
+    case 9: // 512
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 1,16,32,1>;
+        break;
+    case 10: // 1024
+        kernel = &time_masked_softmax_warp_forward<input_t, output_t, acc_t, 1,32,32,1>;
+        break;
+    default:
+        return false;
+    }
+    return true;
+}
+ 
+template<typename input_t, typename output_t, typename acc_t>
+bool dispatch_time_masked_softmax(output_t *dst, const input_t *src, const uint8_t *pad_mask, int softmax_elements, int softmax_elements_stride, int batch_count, int mod_seq_len)
+{
+    if (softmax_elements == 0) {
+        return true;
+    } else if (softmax_elements <= 1024) {
+        // compute function index. there's a function for each power of two size up to 1024.
+        int log2_elements = 0;
+        while ((1 << log2_elements) < softmax_elements) ++log2_elements;
+ 
+        time_masked_softmax_forward_func<input_t, output_t> kernel;
+        int warp_size, batches_per_warp;
+        if (!warp_time_masked_softmax_kernel<input_t, output_t, acc_t>(log2_elements, warp_size, batches_per_warp, kernel)) {
+            return false;
+        }
+ 
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+ 
+        // compute warps per block.
+        int warps_per_block = (threads_per_block / warp_size);
+ 
+        // compute launch size
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+ 
+        // launch
+        kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, pad_mask, batch_count, softmax_elements_stride, softmax_elements, mod_seq_len);
+        return true;
+    }
+    return false;
+}
+
+int log2_ceil_native(int value) {
+    int log2_value = 0;
+    while ((1 << log2_value) < value) ++log2_value;
+    return log2_value;
+}
+
+template <typename T>
+__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff)
+{
+#if CUDA_VERSION >= 9000
+    return __shfl_xor_sync(mask, value, laneMask, width);
+#else
+    return __shfl_xor(value, laneMask, width);
+#endif
+}
+
+template <typename acc_t, int WARP_BATCH, int WARP_SIZE>
+__device__ __forceinline__ void warp_reduce_sum(acc_t* sum) {
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;  i < WARP_BATCH;  ++i) {
+            acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);
+            sum[i] = sum[i] + b;
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+// Warp softmax backward functions as fused variants of at::softmax_backward_data function
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+//softmax backward data function is taken from native pytorch, elementwise mul is fused in the epolog, as well as masking and scaling for fusing dropout
+
+template <typename input_t, typename output_t, typename acc_t, int log2_elements, bool is_log_softmax>
+__global__ void masked_scale_softmax_warp_backward_masked_dgrad(output_t *gradInput, const input_t *grad, const input_t *output, const uint8_t *mask, const uint8_t *pad_mask, acc_t scale, int batch_size, int stride, int element_count, int heads)
+{
+    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method warp_softmax_backward_kernel.
+    constexpr int next_power_of_two = 1 << log2_elements;
+    constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+    constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
+    constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
+
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x % WARP_SIZE;
+
+    // the first element to process by the current thread
+    int thread_offset = first_batch * stride + local_idx;
+    grad += thread_offset;
+    output += thread_offset;
+    gradInput += thread_offset;
+    mask += thread_offset;
+
+    // The nested loops over WARP_BATCH and then WARP_ITERATIONS can be simplified to one loop,
+    // but I think doing so would obfuscate the logic of the algorithm, thus I chose to keep
+    // the nested loops.
+    // This should have no impact on performance because the loops are unrolled anyway.
+
+    // load data from global memory
+    acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS]  ;
+    acc_t output_reg[WARP_BATCH][WARP_ITERATIONS] ;
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        for (int it = 0;  it < WARP_ITERATIONS;  ++it) {
+            int element_index = local_idx + it * WARP_SIZE;
+            if (element_index < batch_element_count) {
+                grad_reg[i][it] = (input_t)((acc_t)mask[i*element_count+it*WARP_SIZE] * (acc_t)grad[i*element_count+it*WARP_SIZE] * (acc_t)scale )*output[i*element_count+it*WARP_SIZE];
+                output_reg[i][it] = output[i*element_count+it*WARP_SIZE];
+            } else {
+                grad_reg[i][it] = acc_t(0);
+                output_reg[i][it] = acc_t(0);
+            }
+        }
+    }
+
+    acc_t sum[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        sum[i] = grad_reg[i][0]; 
+        #pragma unroll
+        for (int it = 1;  it < WARP_ITERATIONS;  ++it) {
+            sum[i] += grad_reg[i][it];
+        }
+    }
+    warp_reduce_sum<acc_t, WARP_BATCH, WARP_SIZE>(sum);
+
+    // store result
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  ++it) {
+            int element_index = local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                // compute gradients
+		int total_ind = thread_offset + i*element_count + it*WARP_SIZE;
+		int pad_mask_ind =  element_count*(total_ind/(heads * element_count * element_count)) + total_ind%element_count;
+		uint8_t pad_mask_element = 1 - pad_mask[pad_mask_ind];
+		if (pad_mask_element == 0) gradInput[i*element_count+it*WARP_SIZE] = 0;
+		else {
+                  if (is_log_softmax) {
+                    gradInput[i*element_count+it*WARP_SIZE] = (grad_reg[i][it] - std::exp(output_reg[i][it]) * sum[i]);
+                  } else {
+                    gradInput[i*element_count+it*WARP_SIZE] = (grad_reg[i][it] - output_reg[i][it] * sum[i]);
+                  }
+		}
+            }
+        }
+    }
+}
+template<typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
+void dispatch_masked_scale_softmax_backward_masked_out(output_t *grad_input, const input_t *grad, const input_t *output, const uint8_t *mask, const uint8_t *pad_mask, acc_t scale, int softmax_elements, int softmax_elements_stride, int batch_count, int heads)
+{
+    TORCH_INTERNAL_ASSERT( softmax_elements >= 0 && softmax_elements <= 1024 );
+    if (softmax_elements == 0) {
+       return;
+    } else {
+        int log2_elements = log2_ceil_native(softmax_elements);
+        const int next_power_of_two = 1 << log2_elements;
+
+        // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_backward.
+        int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+
+        // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_backward.
+        int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+
+        int warps_per_block = (threads_per_block / warp_size);
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+        // Launch code would be more elegant if C++ supported FOR CONSTEXPR
+        switch (log2_elements) {
+            case 0: // 1
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 0, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 1: // 2
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 1, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 2: // 4
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 2, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 3: // 8
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 3, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 4: // 16
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 4, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 5: // 32
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 5, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 6: // 64
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 6, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 7: // 128
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 7, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 8: // 256
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 8, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 9: // 512
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 9, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 10: // 1024
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 10, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            default:
+                break;
+        }
+    }
+}
+
+template<typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
+void dispatch_masked_scale_softmax_backward_masked_out_stream(output_t *grad_input, const input_t *grad, const input_t *output, const uint8_t *mask, const uint8_t *pad_mask, acc_t scale, int softmax_elements, int softmax_elements_stride, int batch_count, int heads, cudaStream_t streamid)
+{
+    TORCH_INTERNAL_ASSERT( softmax_elements >= 0 && softmax_elements <= 1024 );
+    if (softmax_elements == 0) {
+       return;
+    } else {
+        int log2_elements = log2_ceil_native(softmax_elements);
+        const int next_power_of_two = 1 << log2_elements;
+        // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_backward.
+        int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+        // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_backward.
+        int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+        int warps_per_block = (threads_per_block / warp_size);
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+        // Launch code would be more elegant if C++ supported FOR CONSTEXPR
+        switch (log2_elements) {
+            case 0: // 1
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 0, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 1: // 2
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 1, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 2: // 4
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 2, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 3: // 8
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 3, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 4: // 16
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 4, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 5: // 32
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 5, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 6: // 64
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 6, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 7: // 128
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 7, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 8: // 256
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 8, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 9: // 512
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 9, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            case 10: // 1024
+                masked_scale_softmax_warp_backward_masked_dgrad<input_t, output_t, acc_t, 10, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, pad_mask, scale, batch_count, softmax_elements_stride, softmax_elements, heads);
+                break;
+            default:
+                break;
+        }
+    }
+}
+
+template <typename input_t, typename output_t, typename acc_t, int log2_elements, bool is_log_softmax>
+__global__ void masked_scale_softmax_warp_backward(output_t *gradInput, const input_t *grad, const input_t *output, const uint8_t *mask, acc_t scale, int batch_size, int stride, int element_count)
+{
+    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method warp_softmax_backward_kernel.
+    constexpr int next_power_of_two = 1 << log2_elements;
+    constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+    constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
+    constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
+
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x % WARP_SIZE;
+
+    // the first element to process by the current thread
+    int thread_offset = first_batch * stride + local_idx;
+    grad += thread_offset;
+    output += thread_offset;
+    gradInput += thread_offset;
+    mask += thread_offset;
+
+    // The nested loops over WARP_BATCH and then WARP_ITERATIONS can be simplified to one loop,
+    // but I think doing so would obfuscate the logic of the algorithm, thus I chose to keep
+    // the nested loops.
+    // This should have no impact on performance because the loops are unrolled anyway.
+
+    // load data from global memory
+    acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS]  ;
+    acc_t output_reg[WARP_BATCH][WARP_ITERATIONS] ;
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        for (int it = 0;  it < WARP_ITERATIONS;  ++it) {
+            int element_index = local_idx + it * WARP_SIZE;
+            if (element_index < batch_element_count) {
+                grad_reg[i][it] = (input_t)((acc_t)mask[i*element_count+it*WARP_SIZE] * (acc_t)grad[i*element_count+it*WARP_SIZE] * (acc_t)scale )*output[i*element_count+it*WARP_SIZE];
+                output_reg[i][it] = output[i*element_count+it*WARP_SIZE];
+            } else {
+                grad_reg[i][it] = acc_t(0);
+                output_reg[i][it] = acc_t(0);
+            }
+        }
+    }
+
+    acc_t sum[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        sum[i] = grad_reg[i][0]; 
+        #pragma unroll
+        for (int it = 1;  it < WARP_ITERATIONS;  ++it) {
+            sum[i] += grad_reg[i][it];
+        }
+    }
+    warp_reduce_sum<acc_t, WARP_BATCH, WARP_SIZE>(sum);
+
+    // store result
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  ++it) {
+            int element_index = local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                // compute gradients
+                if (is_log_softmax) {
+                    gradInput[i*element_count+it*WARP_SIZE] = (grad_reg[i][it] - std::exp(output_reg[i][it]) * sum[i]);
+                } else {
+                    gradInput[i*element_count+it*WARP_SIZE] = (grad_reg[i][it] - output_reg[i][it] * sum[i]);
+                }
+            }
+        }
+    }
+}
+
+
+
+
+
+template <typename input_t, typename output_t, typename acc_t, int WARP_BATCH, int WARP_ITERATIONS, int WARP_SIZE=32, int ELEMENTS_PER_LDG_STG, bool is_log_softmax>
+__global__ void masked_scale_softmax_warp_backward_recompute(output_t *gradInput, const input_t *grad, const input_t *softmax_input, const input_t *pad_mask, const uint8_t *mask, acc_t scale, int batch_size, int stride, int pad_batch_stride, int element_count)
+{
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x % WARP_SIZE;
+    //vectorize if a row length is multiple of 4
+    int flag_vec4 = element_count & 3 == 0;
+    acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS]  ;
+    input_t elements_input[WARP_BATCH][WARP_ITERATIONS] ;
+
+    // the first element to process by the current thread
+    int thread_offset =  first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+    
+    grad += thread_offset;
+    softmax_input += thread_offset;
+    gradInput += thread_offset;
+    mask += thread_offset;
+    
+    // The nested loops over WARP_BATCH and then WARP_ITERATIONS can be simplified to one loop,
+    // but I think doing so would obfuscate the logic of the algorithm, thus I chose to keep
+    // the nested loops.
+    // This should have no impact on performance because the loops are unrolled anyway.
+    
+    // load data from global memory
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        int pad_thread_offset = ( (first_batch + i) / pad_batch_stride) * stride + ELEMENTS_PER_LDG_STG * local_idx;
+        const input_t* curr_mask    = pad_mask + pad_thread_offset;
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+    
+            #pragma unroll
+            for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+    	//masking_value is a large negative value
+                elements_input[i][it + element] = -10000;
+    	        grad_reg[i][it+element] = acc_t(0);
+            }
+    
+            if (element_index < batch_element_count) {
+                int itr_jmp = it * WARP_SIZE;
+                int itr_idx = i * element_count + itr_jmp;
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], softmax_input + itr_idx);
+                apply_additive_mask<input_t, ELEMENTS_PER_LDG_STG>(&elements_input[i][it], curr_mask + itr_jmp); //(__half)-std::numeric_limits<float>::infinity()
+                uint8_t mask_temp[ELEMENTS_PER_LDG_STG];
+                input_t grad_temp[ELEMENTS_PER_LDG_STG];
+                copy_vector<uint8_t, ELEMENTS_PER_LDG_STG>(&mask_temp[0], mask + itr_idx);
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(&grad_temp[0], grad + itr_idx);
+                #pragma unroll
+                for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                    grad_reg[i][it+element] = ((acc_t)mask_temp[element] * (acc_t)grad_temp[element] * (acc_t)scale );
+                }
+            } 
+    
+        }
+    }
+    // load data from global memory
+ 
+    // convert input_t to acc_t
+    // TODO : remove this, input is already acc_t type in register
+    acc_t elements[WARP_BATCH][WARP_ITERATIONS] ;
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            elements[i][it] = elements_input[i][it];
+        }
+    }
+ 
+    constexpr uint32_t  FULL_MASK = 0xffffffff;
+ 
+    // compute local max_value
+ 
+    // take the max_value of the first element to avoid one max call
+    acc_t max_value[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        max_value[i] = elements[i][0];
+    }
+ 
+    #pragma unroll
+    for (int it = 1;it < WARP_ITERATIONS;++it) {
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
+        }
+    }
+ 
+    // reduction max_value
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        float val[WARP_BATCH];
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            val[i] = __shfl_xor_sync(FULL_MASK, max_value[i], offset, WARP_SIZE);
+        }
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            max_value[i] = max_value[i] > val[i] ? max_value[i] : val[i];
+        }
+    }
+ 
+    // compute local sum
+    acc_t sum[WARP_BATCH] { 0.0f };
+ 
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            //elements[i][it] = expf(elements[i][it] - max_value[i]);
+            elements[i][it] = std::exp(elements[i][it] - max_value[i]);
+            sum[i] += elements[i][it];
+        }
+    }
+ 
+    // reduction sum
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);
+        }
+    }
+
+    // store result
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it ++) {
+	   elements[i][it] = elements[i][it] / sum[i]; 
+           grad_reg[i][it] = grad_reg[i][it] * elements[i][it];
+	}
+    }
+
+    acc_t grad_sum[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        grad_sum[i] = grad_reg[i][0]; 
+        #pragma unroll
+        for (int it = 1;  it < WARP_ITERATIONS;  ++it) {
+            grad_sum[i] += grad_reg[i][it];
+        }
+    }
+    warp_reduce_sum<acc_t, WARP_BATCH, WARP_SIZE>(grad_sum);
+
+    // store result
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                // compute gradients
+	            output_t grad_input_reg[ELEMENTS_PER_LDG_STG];
+                #pragma unroll
+	            for (int element=0; element<ELEMENTS_PER_LDG_STG; element++) {
+                    if (is_log_softmax) {
+                        grad_input_reg[element] = (grad_reg[i][it+element] - std::exp(elements[i][it+element]) * grad_sum[i]);
+                    } else {
+                        grad_input_reg[element] = (grad_reg[i][it+element] - elements[i][it+element] * grad_sum[i]);
+                    }
+	             
+	            }
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(gradInput + i * element_count + it * WARP_SIZE, grad_input_reg);
+            }
+        }
+    }
+}
+
+
+
+template<typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
+using masked_scale_softmax_warp_backward_recompute_func = void(*)(output_t *gradInput, const input_t *grad, const input_t *softmax_input, const input_t *pad_mask, const uint8_t *mask, acc_t scale, int batch_size, int stride, int pad_batch_stride, int element_count);
+
+template <typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
+bool masked_scale_softmax_warp_backward_recompute_kernel(int element_count, int log2_elements, int &warp_size, int &batches_per_warp, masked_scale_softmax_warp_backward_recompute_func<input_t, output_t, acc_t, is_log_softmax> &kernel) {
+    // determine size of a warp
+    const int next_power_of_two = 1 << log2_elements;
+    warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;
+ 
+    // determine how many batches a warp should process.
+    batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+    bool flag_vec4 = (element_count % 4 == 0); 
+    switch (log2_elements) {
+    case 0: // 1
+        kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 2,1,1,1, is_log_softmax>;
+        break;
+    case 1: // 2
+        kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 2,1,2,1, is_log_softmax>;
+        break;
+    case 2: // 4
+        kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 2,1,4,1, is_log_softmax>;
+        break;
+    case 3: // 8
+        kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 2,1,8,1, is_log_softmax>;
+        break;
+    case 4: // 16
+        kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 2,1,16,1, is_log_softmax>;
+        break;
+    case 5: // 32
+        kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 2,1,32,1, is_log_softmax>;
+        break;
+    case 6: // 64
+        kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 2,2,32,1, is_log_softmax>;
+        break;
+    case 7: // 128
+        kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 2,4,32,1, is_log_softmax>;
+        break;
+    case 8: // 256
+	if (flag_vec4) kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 1,8,32,4, is_log_softmax>;
+	else kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 1,8,32,1, is_log_softmax>;
+        break;
+    case 9: // 512
+        if (flag_vec4) kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 1,16,32,4, is_log_softmax>;
+	else kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 1,16,32,1, is_log_softmax>;
+        break;
+    case 10: // 1024
+        if (flag_vec4) kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 1,32,32,4, is_log_softmax>;
+	else kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 1,32,32,1, is_log_softmax>;
+        break;
+    case 11: // 2048
+        if (flag_vec4) kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 1,64,32,4, is_log_softmax>;
+	else kernel = &masked_scale_softmax_warp_backward_recompute<input_t, output_t, acc_t, 1,64,32,1, is_log_softmax>;
+        break;
+    default:
+        return false;
+    }
+    return true;
+}
+
+template<typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
+bool dispatch_masked_scale_softmax_backward_recompute(output_t *grad_input, const input_t *grad, const input_t *softmax_input, const input_t *pad_mask, const uint8_t *mask, acc_t scale, int softmax_elements, int softmax_elements_stride, int pad_batch_stride, int batch_count, cudaStream_t streamid)
+{
+	
+    if (softmax_elements == 0) {
+        return true;
+    } else if (softmax_elements <= 2048) {
+        // compute function index. there's a function for each power of two size up to 1024.
+        int log2_elements = 0;
+        while ((1 << log2_elements) < softmax_elements) ++log2_elements;
+ 
+        masked_scale_softmax_warp_backward_recompute_func<input_t, output_t, acc_t, is_log_softmax> kernel;
+        int warp_size, batches_per_warp;
+        if (!masked_scale_softmax_warp_backward_recompute_kernel<input_t, output_t, acc_t, is_log_softmax>(softmax_elements, log2_elements, warp_size, batches_per_warp, kernel)) {
+            return false;
+        }
+ 
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+        // compute warps per block.
+        int warps_per_block = (threads_per_block / warp_size);
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+
+        // compute launch size
+        dim3 threads(warp_size, warps_per_block, 1);
+         
+        // launch
+        kernel<<<blocks, threads, 0, streamid>>>(grad_input, grad, softmax_input, pad_mask, mask, scale, batch_count, softmax_elements_stride, pad_batch_stride, softmax_elements);
+        return true;
+    }
+    return false;
+}
+
+
+template<typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
+void dispatch_masked_scale_softmax_backward_stream(output_t *grad_input, const input_t *grad, const input_t *output, const uint8_t *mask, acc_t scale, int softmax_elements, int softmax_elements_stride, int batch_count, cudaStream_t streamid)
+{
+    TORCH_INTERNAL_ASSERT( softmax_elements >= 0 && softmax_elements <= 1024 );
+    if (softmax_elements == 0) {
+       return;
+    } else {
+        int log2_elements = log2_ceil_native(softmax_elements);
+        const int next_power_of_two = 1 << log2_elements;
+        // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_backward.
+        int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+        // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_backward.
+        int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+        int warps_per_block = (threads_per_block / warp_size);
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+        // Launch code would be more elegant if C++ supported FOR CONSTEXPR
+        switch (log2_elements) {
+            case 0: // 1
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 0, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 1: // 2
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 1, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 2: // 4
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 2, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 3: // 8
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 3, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 4: // 16
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 4, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 5: // 32
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 5, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 6: // 64
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 6, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 7: // 128
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 7, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 8: // 256
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 8, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 9: // 512
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 9, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 10: // 1024
+                masked_scale_softmax_warp_backward<input_t, output_t, acc_t, 10, is_log_softmax>
+                    <<<blocks, threads, 0, streamid>>>(grad_input, grad, output, mask, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            default:
+                break;
+        }
+    }
+}
+
+// elementwise multiplication called in at::softmax_backward_data is fused inside softmax dgrad kernel
+// as a result of fusion, intermediate multiplication result is stored in fp32 in registers, instead of fp16
+template <typename input_t, typename output_t, typename acc_t, int log2_elements, bool is_log_softmax>
+__global__ void softmax_warp_backward_fused_native(output_t *gradInput, const input_t *grad, const input_t *output, int batch_size, int stride, int element_count)
+{
+    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and warp_size of method warp_softmax_backward_kernel.
+    constexpr int next_power_of_two = 1 << log2_elements;
+    constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+    constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
+    constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
+
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x % WARP_SIZE;
+
+    // the first element to process by the current thread
+    int thread_offset = first_batch * stride + local_idx;
+    grad += thread_offset;
+    output += thread_offset;
+    gradInput += thread_offset;
+
+    // The nested loops over WARP_BATCH and then WARP_ITERATIONS can be simplified to one loop,
+    // but I think doing so would obfuscate the logic of the algorithm, thus I chose to keep
+    // the nested loops.
+    // This should have no impact on performance because the loops are unrolled anyway.
+
+    // load data from global memory
+    acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS]  ;
+    acc_t output_reg[WARP_BATCH][WARP_ITERATIONS] ;
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        for (int it = 0;  it < WARP_ITERATIONS;  ++it) {
+            int element_index = local_idx + it * WARP_SIZE;
+            if (element_index < batch_element_count) {
+                grad_reg[i][it] = grad[i*element_count+it*WARP_SIZE]*output[i*element_count+it*WARP_SIZE];
+                output_reg[i][it] = output[i*element_count+it*WARP_SIZE];
+            } else {
+                grad_reg[i][it] = acc_t(0);
+                output_reg[i][it] = acc_t(0);
+            }
+        }
+    }
+
+    acc_t sum[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        sum[i] = grad_reg[i][0]; //* output_reg[i][0];
+        #pragma unroll
+        for (int it = 1;  it < WARP_ITERATIONS;  ++it) {
+            sum[i] += grad_reg[i][it];// * output_reg[i][it];
+        }
+    }
+    warp_reduce_sum<acc_t, WARP_BATCH, WARP_SIZE>(sum);
+
+    // store result
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  ++it) {
+            int element_index = local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                // compute gradients
+                if (is_log_softmax) {
+                    gradInput[i*element_count+it*WARP_SIZE] = (grad_reg[i][it] - std::exp(output_reg[i][it]) * sum[i]);
+                } else {
+                    gradInput[i*element_count+it*WARP_SIZE] = (grad_reg[i][it] - output_reg[i][it] * sum[i]);
+                }
+            }
+        }
+    }
+}
+
+template<typename input_t, typename output_t, typename acc_t, bool is_log_softmax>
+void dispatch_softmax_backward_fused_native(output_t *grad_input, const input_t *grad, const input_t *output, int softmax_elements, int softmax_elements_stride, int batch_count)
+{
+    TORCH_INTERNAL_ASSERT( softmax_elements >= 0 && softmax_elements <= 1024 );
+    if (softmax_elements == 0) {
+       return;
+    } else {
+        int log2_elements = log2_ceil_native(softmax_elements);
+        const int next_power_of_two = 1 << log2_elements;
+
+        // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_backward.
+        int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+
+        // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_backward.
+        int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+
+        int warps_per_block = (threads_per_block / warp_size);
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+        // Launch code would be more elegant if C++ supported FOR CONSTEXPR
+        switch (log2_elements) {
+            case 0: // 1
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 0, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 1: // 2
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 1, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 2: // 4
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 2, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 3: // 8
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 3, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 4: // 16
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 4, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 5: // 32
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 5, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 6: // 64
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 6, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 7: // 128
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 7, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 8: // 256
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 8, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 9: // 512
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 9, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 10: // 1024
+                softmax_warp_backward_fused_native<input_t, output_t, acc_t, 10, is_log_softmax>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            default:
+                break;
+        }
+    }
+}
+
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+// Warp softmax backward
+////////////////////////////////////////////////////////////////////////////////////////////////////////////////////////
+
+template <typename input_t, typename output_t, typename acc_t, int WARP_BATCH, int WARP_ITERATIONS, int WARP_SIZE=32, int ELEMENTS_PER_LDG_STG=1>
+__global__ void softmax_warp_backward(__half *gradInput, const __half *grad, const __half *output, int batch_size, int stride, int element_count)
+{
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+ 
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+ 
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+ 
+    // the first element to process by the current thread
+    int thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+    grad += thread_offset;
+    output += thread_offset;
+    gradInput += thread_offset;
+ 
+    // load data from global memory
+    input_t grad_reg_input[WARP_BATCH][WARP_ITERATIONS] = {0.0f};
+    input_t output_reg_input[WARP_BATCH][WARP_ITERATIONS] = {0.0f};
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < batch_element_count) {
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(&grad_reg_input[i][it], grad + i * element_count + it * WARP_SIZE);
+                copy_vector<input_t,ELEMENTS_PER_LDG_STG>(&output_reg_input[i][it], output + i * element_count + it * WARP_SIZE);
+            }
+ 
+        }
+    }
+ 
+    // convert half to floating point
+    acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS];
+    acc_t output_reg[WARP_BATCH][WARP_ITERATIONS];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            grad_reg[i][it] = grad_reg_input[i][it];
+            output_reg[i][it] = output_reg_input[i][it];
+        }
+    }
+ 
+ 
+    // compute thread local sum
+    acc_t sum[WARP_BATCH] = {0};
+    #pragma unroll
+    for (int it = 0;it < WARP_ITERATIONS;++it) {
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += grad_reg[i][it] * output_reg[i][it];
+ 
+        }
+    }
+ 
+    // reduction sum
+    constexpr uint32_t FULL_MASK = 0xffffffff;
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);
+        }
+    }
+ 
+    // store result
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                // compute gradients
+                output_t out[ELEMENTS_PER_LDG_STG];
+                for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                    out[element] = (output_reg[i][it+element] * (grad_reg[i][it+element] - sum[i]));
+                }
+                // store them in global memory
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(gradInput + i * element_count + it * WARP_SIZE, out);
+            }
+        }
+    }
+}
+ 
+ 
+ 
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t>
+using softmax_backward_func = void(*)(output_t *gradInput, const input_t *grad, const input_t *output, int batch_size, int stride, int element_count);
+ 
+template <typename input_t, typename output_t, typename acc_t>
+bool warp_softmax_backward_kernel(int log2_elements, int &warp_size, int &batches_per_warp, softmax_backward_func<input_t, output_t> &kernel) {
+    // determine size of a warp
+    const int next_power_of_two = 1 << log2_elements;
+    warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;
+ 
+    // determine how many batches a warp should process.
+    batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+ 
+    switch (log2_elements) {
+    case 0: // 1
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 2,1,1,1>;
+        break;
+    case 1: // 2
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 2,1,2,1>;
+        break;
+    case 2: // 4
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 2,1,4,1>;
+        break;
+    case 3: // 8
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 2,1,8,1>;
+        break;
+    case 4: // 16
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 2,1,16,1>;
+        break;
+    case 5: // 32
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 2,1,32,1>;
+        break;
+    case 6: // 64
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 2,2,32,1>;
+        break;
+    case 7: // 128
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 2,4,32,1>;
+        break;
+    case 8: // 256
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 1,8,32,1>;
+        break;
+    case 9: // 512
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 1,16,32,1>;
+        break;
+    case 10: // 1024
+        kernel = &softmax_warp_backward<input_t, output_t, acc_t, 1,32,32,1>;
+        break;
+    default:
+        return false;
+    }
+    return true;
+}
+ 
+template<typename input_t, typename output_t, typename acc_t>
+bool dispatch_softmax_backward(output_t *grad_input, const input_t *grad, const input_t *output, int softmax_elements, int softmax_elements_stride, int batch_count)
+{
+    if (softmax_elements == 0) {
+        return true;
+    } else if (softmax_elements <= 1024) {
+        // compute function index. there's a function for each power of two size up to 1024.
+        int log2_elements = 0;
+        while ((1 << log2_elements) < softmax_elements) ++log2_elements;
+ 
+        softmax_backward_func<input_t, output_t> kernel;
+        int warp_size, batches_per_warp;
+        if (!warp_softmax_backward_kernel<input_t, output_t, acc_t>(log2_elements, warp_size, batches_per_warp, kernel)) {
+            return false;
+        }
+ 
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+ 
+        // compute warps per block.
+        int warps_per_block = (threads_per_block / warp_size);
+ 
+        // compute launch size
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+ 
+        // launch
+        kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+        return true;
+    }
+    return false;
+}
+
+template<typename input_t, typename output_t, typename acc_t>
+bool dispatch_softmax_backward_stream(output_t *grad_input, const input_t *grad, const input_t *output, int softmax_elements, int softmax_elements_stride, int batch_count, cudaStream_t streamid)
+{
+    if (softmax_elements == 0) {
+        return true;
+    } else if (softmax_elements <= 1024) {
+        // compute function index. there's a function for each power of two size up to 1024.
+        int log2_elements = 0;
+        while ((1 << log2_elements) < softmax_elements) ++log2_elements;
+        softmax_backward_func<input_t, output_t> kernel;
+        int warp_size, batches_per_warp;
+        if (!warp_softmax_backward_kernel<input_t, output_t, acc_t>(log2_elements, warp_size, batches_per_warp, kernel)) {
+            return false;
+        }
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+        // compute warps per block.
+        int warps_per_block = (threads_per_block / warp_size);
+        // compute launch size
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+        // launch
+        kernel<<<blocks, threads, 0, streamid>>>(grad_input, grad, output, batch_count, softmax_elements_stride, softmax_elements);
+        return true;
+    }
+    return false;
+}
+
+template <typename input_t, typename output_t, typename acc_t, int WARP_BATCH, int WARP_ITERATIONS, int WARP_SIZE=32, int ELEMENTS_PER_LDG_STG=1>
+__global__ void masked_softmax_warp_backward(__half *gradInput, const __half *grad, const __half *output, const uint8_t *pad_mask, int batch_size, int stride, int element_count, int pad_batch_stride)
+{
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+ 
+    // batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+ 
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+ 
+    // the first element to process by the current thread
+    int thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+    grad += thread_offset;
+    output += thread_offset;
+    gradInput += thread_offset;
+ 
+    // load data from global memory
+    input_t grad_reg_input[WARP_BATCH][WARP_ITERATIONS] = {0.0f};
+    input_t output_reg_input[WARP_BATCH][WARP_ITERATIONS] = {0.0f};
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < batch_element_count) {
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(&grad_reg_input[i][it], grad + i * element_count + it * WARP_SIZE);
+                copy_vector<input_t,ELEMENTS_PER_LDG_STG>(&output_reg_input[i][it], output + i * element_count + it * WARP_SIZE);
+            }
+ 
+        }
+    }
+ 
+    // convert half to floating point
+    acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS];
+    acc_t output_reg[WARP_BATCH][WARP_ITERATIONS];
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        for (int it = 0;it < WARP_ITERATIONS;++it) {
+            grad_reg[i][it] = grad_reg_input[i][it];
+            output_reg[i][it] = output_reg_input[i][it];
+        }
+    }
+ 
+ 
+    // compute thread local sum
+    acc_t sum[WARP_BATCH] = {0};
+    #pragma unroll
+    for (int it = 0;it < WARP_ITERATIONS;++it) {
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += grad_reg[i][it] * output_reg[i][it];
+ 
+        }
+    }
+ 
+    // reduction sum
+    constexpr uint32_t FULL_MASK = 0xffffffff;
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;i < WARP_BATCH;++i) {
+            sum[i] += __shfl_xor_sync(FULL_MASK, sum[i], offset, WARP_SIZE);
+        }
+    }
+ 
+    // store result
+    #pragma unroll
+    for (int i = 0;i < WARP_BATCH;++i) {
+        if (i >= local_batches)
+            break;
+        int pad_thread_offset = ( (first_batch + i) / pad_batch_stride) * stride + ELEMENTS_PER_LDG_STG * local_idx;
+        const uint8_t* curr_mask    = pad_mask + pad_thread_offset;
+        #pragma unroll
+        for (int it = 0;it < WARP_ITERATIONS;it += ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                // compute gradients
+                output_t out[ELEMENTS_PER_LDG_STG];
+                for (int element = 0;element < ELEMENTS_PER_LDG_STG;++element) {
+                    out[element] = (output_reg[i][it+element] * (grad_reg[i][it+element] - sum[i]));
+                }
+                // store them in global memory
+                int itr_jmp = it * WARP_SIZE;
+                int itr_idx = i * element_count + itr_jmp;
+                // It is kind of unfortunate this has to be here to zero something out that is close to
+                // zero in the first place
+                apply_mask<input_t, ELEMENTS_PER_LDG_STG>(&out[0], 0.0, curr_mask + itr_jmp);
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(gradInput + itr_idx, out);
+            }
+        }
+    }
+}
+ 
+ 
+ 
+// WARP_BATCH number of batches.
+// WARP_ITERATOINS The number of iterations required for one warp to iterate over all data.
+// WARP_SIZE number of elements working on a single batch, has to be a power of two.
+// ELEMENTS_PER_LDG_STG has to be 1.
+template <typename input_t, typename output_t>
+using masked_softmax_backward_func = void(*)(output_t *gradInput, const input_t *grad, const input_t *output, const uint8_t *pad_mask, int batch_size, int stride, int element_count, int pad_batch_stride);
+ 
+template <typename input_t, typename output_t, typename acc_t>
+bool warp_masked_softmax_backward_kernel(int log2_elements, int &warp_size, int &batches_per_warp, masked_softmax_backward_func<input_t, output_t> &kernel) {
+    // determine size of a warp
+    const int next_power_of_two = 1 << log2_elements;
+    warp_size = (next_power_of_two < 32) ? next_power_of_two : 32;
+ 
+    // determine how many batches a warp should process.
+    batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+ 
+    switch (log2_elements) {
+    case 0: // 1
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 2,1,1,1>;
+        break;
+    case 1: // 2
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 2,1,2,1>;
+        break;
+    case 2: // 4
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 2,1,4,1>;
+        break;
+    case 3: // 8
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 2,1,8,1>;
+        break;
+    case 4: // 16
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 2,1,16,1>;
+        break;
+    case 5: // 32
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 2,1,32,1>;
+        break;
+    case 6: // 64
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 2,2,32,1>;
+        break;
+    case 7: // 128
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 2,4,32,1>;
+        break;
+    case 8: // 256
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 1,8,32,1>;
+        break;
+    case 9: // 512
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 1,16,32,1>;
+        break;
+    case 10: // 1024
+        kernel = &masked_softmax_warp_backward<input_t, output_t, acc_t, 1,32,32,1>;
+        break;
+    default:
+        return false;
+    }
+    return true;
+}
+ 
+template<typename input_t, typename output_t, typename acc_t>
+bool dispatch_masked_softmax_backward(output_t *grad_input, const input_t *grad, const input_t *output, const uint8_t *pad_mask, int softmax_elements, int softmax_elements_stride, int batch_count, int pad_batch_stride)
+{
+    if (softmax_elements == 0) {
+        return true;
+    } else if (softmax_elements <= 1024) {
+        // compute function index. there's a function for each power of two size up to 1024.
+        int log2_elements = 0;
+        while ((1 << log2_elements) < softmax_elements) ++log2_elements;
+ 
+        masked_softmax_backward_func<input_t, output_t> kernel;
+        int warp_size, batches_per_warp;
+        if (!warp_masked_softmax_backward_kernel<input_t, output_t, acc_t>(log2_elements, warp_size, batches_per_warp, kernel)) {
+            return false;
+        }
+ 
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+ 
+        // compute warps per block.
+        int warps_per_block = (threads_per_block / warp_size);
+ 
+        // compute launch size
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = (batch_count + batches_per_block - 1) / batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+ 
+        // launch
+        kernel<<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, pad_mask, batch_count, softmax_elements_stride, softmax_elements, pad_batch_stride);
+        return true;
+    }
+    return false;
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/strided_batched_gemm.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/strided_batched_gemm.h
new file mode 100644
index 0000000000..26a6337911
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/multihead_attn/strided_batched_gemm.h
@@ -0,0 +1,407 @@
+#include <vector>
+#include <iostream>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+
+//#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+
+#include "cutlass/cutlass.h"
+#include "cutlass/gemm/gemm.h"
+#include "cutlass/gemm/wmma_gemm_traits.h"
+
+// symbol to be automatically resolved by PyTorch libs
+extern THCState *state;
+
+cublasOperation_t convertTransToCublasOperation(char trans) {
+  if (trans == 't') return CUBLAS_OP_T;
+  else if (trans == 'n') return CUBLAS_OP_N;
+  else if (trans == 'c') return CUBLAS_OP_C;
+  else {
+    AT_ERROR("trans must be one of: t, n, c");
+    return CUBLAS_OP_T;
+  }
+}
+
+void CublasStridedBatchedGemm(THCState *state, char transa, char transb, long m, long n, long k,
+                    float alpha, const half *a, long lda, long strideA, const half *b, long ldb, long strideB,
+                    float beta, half *c, long ldc, long strideC, long batchCount, cublasGemmAlgo_t algo=CUBLAS_GEMM_DEFAULT_TENSOR_OP) {
+    cublasOperation_t opa = convertTransToCublasOperation(transa);
+    cublasOperation_t opb = convertTransToCublasOperation(transb);
+
+    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+    cudaStream_t   stream = at::cuda::getCurrentCUDAStream().stream();
+    cublasSetStream(handle, stream);
+    float fAlpha = alpha;
+    float fBeta = beta;
+    //THCublasCheck(cublasSetMathMode(handle, CUBLAS_TENSOR_OP_MATH));
+    TORCH_CUDABLAS_CHECK(cublasGemmStridedBatchedEx(handle,
+                                     opa, opb, (int)m, (int)n, (int)k,
+                                     (void*)&fAlpha, a, CUDA_R_16F, (int)lda, strideA,
+                                     b, CUDA_R_16F, (int)ldb, strideB,
+                                     (void*)&fBeta, c, CUDA_R_16F, (int)ldc, strideC,
+                                     (int)batchCount, CUDA_R_32F, algo));
+    //THCublasCheck(cublasSetMathMode(handle, CUBLAS_DEFAULT_MATH));
+}
+
+template<cutlass::MatrixLayout::Kind A_LAYOUT, cutlass::MatrixLayout::Kind B_LAYOUT, int SRC_A, int SRC_B, int DST_C>
+void CutlassGemm_FP32Accum(cudaStream_t stream, long m, long n, long k,
+                          float alpha, const half *a, long lda, long strideA, const half *b, long ldb, long strideB,
+                          float beta, half *c, long ldc, long strideC, long batchCount) {
+  //printf("CUTLASS-> %c%c M: %ld N: %ld K: %ld %d%d%d LDA: %ld LDB: %ld LDC: %ld strideA: %ld strideB: %ld strideC: %ld Alpha: %f Beta: %f\n", ((int)A_LAYOUT == 0 ? 'T' : 'N'), ((int)B_LAYOUT ==0 ? 'T' : 'N'), m, n, k, SRC_A,SRC_B,DST_C, lda, ldb, ldc, strideA, strideB, strideC, alpha, beta);
+  typedef cutlass::gemm::WmmaGemmTraits<
+    A_LAYOUT,
+    B_LAYOUT,
+    cutlass::Shape<32, 16, 16>,
+    half,
+    half,
+    half,
+    cutlass::gemm::LinearScaling<float>,
+    float,
+    typename cutlass::gemm::WmmaGemmAccumulatorsPerWarp<typename cutlass::Shape<32, 16, 16> >::Shape,
+    typename cutlass::Shape<16, 16, 16>,
+    SRC_A,   //kScalarsPerLdgA_
+    SRC_B,   //kScalarsPerLdgB_
+    SRC_A,   //KScalarsPerLdsA_
+    SRC_B,   //KScalarsPerLdsB_
+    DST_C,   //kScalarsPerLdgCAndStgD_
+    DST_C/2, //kScalarsPerStsD_
+    DST_C/2  //kScalarsPerLdsD_
+  >
+    WmmaGemmTraits;
+
+  typedef cutlass::gemm::Gemm<WmmaGemmTraits> Gemm;
+  typename Gemm::Params params;
+
+
+  int result = params.initialize(
+    m,                  // M dimension for each batch
+    n,                  // N dimension for each batch
+    k,                  // K dimension for each batch
+    alpha,              // scalar alpha
+    a,
+    lda,
+    strideA,     // distance in memory between the first element of neighboring batch
+    b,
+    ldb,
+    strideB,     // distance in memory between the first element of neighboring batch
+    beta,               // scalar beta
+    c,                  // source matrix C
+    ldc,
+    strideC,     // distance in memory between the first element of neighboring batch
+    c,                  // destination matrix C (may be different memory than source C matrix)
+    ldc,
+    strideC,    // distance in memory between the first element of neighboring batch
+    batchCount
+  );
+
+  AT_ASSERTM(result == 0, "Failed to initialize CUTLASS Gemm::Params object.");
+  
+  // batchCount in cutlass batched GEMM kernels maps to gridDim.z, which is limited to 16 bits. 
+  // To implement batched GEMM with larger batch size, we fragment it into
+  // smaller batched GEMMs of gridDim.z <= 64k
+  long batchesLeft    = batchCount;
+  long iterBatchCount = std::min(batchesLeft, static_cast<long>((1 << 16) - 1));
+  
+  do {
+  	 //printf("CUTLASS-> %c%c M: %ld N: %ld K: %ld %d%d%d LDA: %ld LDB: %ld LDC: %ld strideA: %ld strideB: %ld strideC: %ld Alpha: %f Beta: %f TotalBatches: %ld iterBatchCount %ld\n", ((int)A_LAYOUT == 0 ? 'T' : 'N'), ((int)B_LAYOUT ==0 ? 'T' : 'N'), m, n, k, SRC_A,SRC_B,DST_C, lda, ldb, ldc, strideA, strideB, strideC, alpha, beta, batchesLeft, iterBatchCount);
+    int result = params.initialize(
+      m,                  // M dimension for each batch
+      n,                  // N dimension for each batch
+      k,                  // K dimension for each batch
+      alpha,              // scalar alpha
+      a,
+      lda,
+      strideA,     // distance in memory between the first element of neighboring batch
+      b,
+      ldb,
+      strideB,     // distance in memory between the first element of neighboring batch
+      beta,               // scalar beta
+      c,                  // source matrix C
+      ldc,
+      strideC,     // distance in memory between the first element of neighboring batch
+      c,                  // destination matrix C (may be different memory than source C matrix)
+      ldc,
+      strideC,    // distance in memory between the first element of neighboring batch
+      iterBatchCount
+    );
+
+    AT_ASSERTM(result == 0, "Failed to initialize CUTLASS Gemm::Params object.");
+    // Launch the CUTLASS GEMM kernel.
+    C10_CUDA_CHECK(Gemm::launch(params, stream));
+
+    // Update batched GEMM params based on completed work
+    batchesLeft = batchesLeft - iterBatchCount;
+    a += iterBatchCount * strideA;
+    b += iterBatchCount * strideB;
+    c += iterBatchCount * strideC;;
+
+    iterBatchCount = std::min(batchesLeft, static_cast<long>((1 << 16) - 1));
+    
+  } while(batchesLeft > 0);
+}
+
+void gemm_switch_fp32accum(THCState *state, char transa, char transb, long m, long n, long k,
+                           float alpha, const half *a, long lda, long strideA, const half *b, long ldb, long strideB,
+                           float beta, half *c, long ldc, long strideC, long batchCount) {
+  auto stream = c10::cuda::getCurrentCUDAStream();
+  //printf("GEMM   -> %c%c M: %i N: %i K: %i Alpha: %f Beta: %f\n", (transa == 't' ? 'T' : 'N'), (transb =='t' ? 'T' : 'N'), m, n, k, alpha, beta);
+  if        ( (transa == 't') && (transb == 'n') ) { 
+    if      (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x7)) { CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP); }
+    /*if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x7)) {
+      int m_rem = m % 64;
+      int n_rem = n % 64;
+      if ( (m_rem > 48) && ( m <= 192) && (n_rem > 48) && (n <= 192 ) ) {
+        CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP);
+      } else if ( (m_rem > 32) && ( m > 192) && (n_rem > 32) && (n > 192) ) {
+        CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP);
+      } else {
+        CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,8,8,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); 
+      }
+    }*/
+    else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,8,8,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,8,8,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,8,4,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,8,4,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,8,4,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,8,2,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,8,2,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,8,2,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,4,8,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,4,8,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,4,8,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,4,4,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,4,4,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,4,4,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,4,2,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,4,2,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,4,2,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,2,8,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,2,8,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,2,8,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,2,4,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,2,4,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,2,4,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,2,2,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,2,2,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kRowMajor,cutlass::MatrixLayout::kColumnMajor,2,2,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else                                                   { CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+  } else if ( (transa == 'n') && (transb == 'n') ) {
+    if      (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x7)) { CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP); }
+    /*if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x7)) {
+      int m_rem = m % 64;
+      int n_rem = n % 64;
+      if ( (m_rem > 48) && ( m <= 192) && (n_rem > 48) && (n <= 192 ) ) {
+        CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP);
+      } else if ( (m_rem > 32) && ( m > 192) && (n_rem > 32) && (n > 192) ) {
+        CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP);
+      } else {
+        CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,8,8,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);
+      }
+    }*/
+    else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,8,8,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,8,8,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,8,4,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,8,4,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,8,4,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,8,2,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,8,2,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,8,2,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,4,8,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,4,8,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,4,8,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,4,4,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,4,4,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,4,4,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,4,2,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,4,2,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,4,2,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,2,8,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,2,8,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,2,8,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,2,4,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,2,4,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,2,4,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,2,2,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,2,2,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kColumnMajor,2,2,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else                                                   { CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+  } else if ( (transa == 'n') && (transb == 't') ) {
+    if      (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x7)) { CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP); }
+    /*if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x7)) { 
+      int m_rem = m % 64;
+      int n_rem = n % 64;
+      if ( (m_rem > 48) && ( m <= 192) && (n_rem > 48) && (n <= 192 ) ) {
+        CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP); 
+      } else if ( (m_rem > 32) && ( m > 192) && (n_rem > 32) && (n > 192) ) {
+        CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount, CUBLAS_GEMM_ALGO0_TENSOR_OP); 
+      } else {
+        CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,8,8,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); 
+      }
+    }*/
+    else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,8,8,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x7) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,8,8,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,8,4,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,8,4,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x3) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,8,4,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,8,2,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,8,2,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x7) && !(ldb & 0x1) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,8,2,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,4,8,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,4,8,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x7) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,4,8,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,4,4,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x3) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,4,4,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,4,2,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,4,2,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x3) && !(ldb & 0x1) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,4,2,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,2,8,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,2,8,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x7) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,2,8,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,2,4,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,2,4,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x3) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,2,4,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x7)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,2,2,8>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x3)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,2,2,4>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else if (!(lda & 0x1) && !(ldb & 0x1) && !(ldc & 0x1)) { CutlassGemm_FP32Accum<cutlass::MatrixLayout::kColumnMajor,cutlass::MatrixLayout::kRowMajor,2,2,2>(stream, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+    else                                                   { CublasStridedBatchedGemm(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount); }
+  } else {
+    AT_ASSERTM(false, "TransA and TransB are invalid");
+  }
+}
+
+void adjustLdLevel3(char transa, char transb, int64_t m, int64_t n, int64_t k, int64_t *lda, int64_t *ldb, int64_t *ldc)
+{
+  int transa_ = ((transa == 't') || (transa == 'T'));
+  int transb_ = ((transb == 't') || (transb == 'T'));
+
+  // Note: leading dimensions generally are checked that they are > 0 and at least as big the result
+  // requires (even if the value won't be used).
+  if(n <= 1)
+    *ldc = std::max<int64_t>(m, 1);
+
+  if(transa_)
+  {
+    if(m <= 1)
+      *lda = std::max<int64_t>(k, 1);
+  }
+  else
+  {
+    if(k <= 1)
+      *lda = std::max<int64_t>(m, 1);
+  }
+
+  if(transb_)
+  {
+    if(k <= 1)
+      *ldb = std::max<int64_t>(n, 1);
+  }
+  else
+  {
+    if(n <= 1)
+      *ldb = std::max<int64_t>(k, 1);
+  }
+
+}
+
+void HgemmStridedBatched(THCState *state, char transa, char transb, long m, long n, long k,
+                             float alpha, const half *a, long lda, long strideA, const half *b, long ldb, long strideB,
+                             float beta, half *c, long ldc, long strideC, long batchCount)
+{
+  if( (m >= INT_MAX) || (n >= INT_MAX) || (k >= INT_MAX) || (lda >= INT_MAX)  || (ldb >= INT_MAX) || (ldc >= INT_MAX) || (batchCount >= INT_MAX) )
+
+  {
+    AT_ERROR("Cublas_SgemmStridedBatched only supports m, n, k, lda, ldb, ldc, batchCount"
+            "with the bound [val] <= %d", INT_MAX);
+  }
+
+  adjustLdLevel3(transa, transb, m, n, k, &lda, &ldb, &ldc);
+
+  //gemm_switch(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);
+  gemm_switch_fp32accum(state, transa, transb, m, n, k, alpha, a, lda, strideA, b, ldb, strideB, beta, c, ldc, strideC, batchCount);
+}
+
+/******
+at::Tensor strided_batched_gemm_cuda(
+    float beta,
+    at::Tensor in_result,
+    float alpha,
+    at::Tensor batch1,
+    at::Tensor batch2) {
+
+  bool transpose_result;
+  char transpose_batch1, transpose_batch2;
+  int64_t lda, ldb, ldc;
+  at::Tensor result, input1, input2;
+  if (in_result.stride(1) == 1)
+  {
+    transpose_result = false;
+    result = in_result;
+    ldc = result.stride(2);
+  }
+  else if (in_result.stride(2) == 1)
+  {
+    transpose_result = true;
+
+    at::Tensor swap = batch2;
+    batch2 = batch1;
+    batch1 = swap;
+
+    result = in_result;
+    ldc = result.stride(1);
+  } else { 
+    AT_ASSERTM(false, "result should be contiguous");
+  }
+
+  if (batch1.stride(transpose_result ? 2 : 1) == 1 &&
+      batch1.stride(transpose_result ? 1 : 2) != 0) {
+    transpose_batch1 = 'n';
+    input1 = batch1;
+    lda = input1.stride(transpose_result ? 1 : 2);
+  } else if (batch1.stride(transpose_result ? 1 : 2) == 1 &&
+             batch1.stride(transpose_result ? 2 : 1) != 0) {
+    transpose_batch1 = 't';
+    input1 = batch1;
+    lda = input1.stride(transpose_result ? 2 : 1);
+  } else {
+    AT_ASSERTM(false, "input1 should be contiguous");
+  }
+
+  if (batch2.stride(transpose_result ? 2 : 1) == 1 &&
+      batch2.stride(transpose_result ? 1 : 2) != 0) {
+    transpose_batch2 = 'n';
+    input2 = batch2;
+    ldb = input2.stride(transpose_result ? 1 : 2);
+  } else if (batch2.stride(transpose_result ? 1 : 2) == 1 &&
+             batch2.stride(transpose_result ? 2 : 1) != 0) {
+    transpose_batch2 = 't';
+    input2 = batch2;
+    ldb = input2.stride(transpose_result ? 2 : 1);
+  } else {
+    AT_ASSERTM(false, "input2 should be contiguous");
+  }
+  int64_t num_batches = result.size(0);
+
+  HgemmStridedBatched(
+      state,
+      transpose_batch1,
+      transpose_batch2,
+      result.size(transpose_result ? 2 : 1),
+      result.size(transpose_result ? 1 : 2),
+      input1.size(transpose_result ? 1 : 2),
+      alpha,
+      static_cast<const half*>(input1.data_ptr()), lda, input1.stride(0),
+      static_cast<const half*>(input2.data_ptr()), ldb, input2.stride(0),
+      beta,
+      static_cast<half*>(result.data_ptr()), ldc, result.stride(0),
+      num_batches);
+
+  return in_result;
+}
+
+***/
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_adam_cuda.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_adam_cuda.cpp
new file mode 100644
index 0000000000..c03c90fe31
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_adam_cuda.cpp
@@ -0,0 +1,86 @@
+#include <torch/extension.h>
+
+// CUDA forward declaration
+void fused_strided_check_finite(at::Tensor & overflow_flag, at::Tensor & p_copy, int stride, int clear_overflow_first);
+
+void fused_adam_cuda(at::Tensor & p, at::Tensor & p_copy, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay);
+void fused_reversible_adam_cuda(at::Tensor & p, at::Tensor & p_copy, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay);
+void fused_maybe_adam_undo_cuda(at::Tensor & overflow_flag, at::Tensor & p, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay);
+
+void fused_adam_cuda_mt(int chunk_size, at::Tensor overflow_flag, std::vector<std::vector<at::Tensor>> tensor_lists, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay);
+
+void maybe_cast_cuda(at::Tensor & overflow_flag, at::Tensor & p_in, at::Tensor & p_out);
+void maybe_cast_cuda_mt(int chunk_size, at::Tensor overflow_flag, std::vector<std::vector<at::Tensor>> tensor_lists);
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+// C++ interface
+void strided_check_finite(
+		at::Tensor& overflow_flag,
+		at::Tensor& p_copy,
+		int stride,
+		int clear_overflow_first
+	 ) {
+	CHECK_INPUT(p_copy);
+	fused_strided_check_finite(overflow_flag, p_copy, stride, clear_overflow_first);
+}
+void adam(at::Tensor & p, at::Tensor & p_copy, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay) {
+        CHECK_INPUT(p);
+        if (p_copy.numel() > 0) CHECK_INPUT(p_copy);
+        CHECK_INPUT(m);
+        CHECK_INPUT(v);
+        CHECK_INPUT(g);
+        int64_t num_elem = p.numel();
+        AT_ASSERTM(m.numel() == num_elem, "number of elements in m and p tensors should be equal");
+        AT_ASSERTM(v.numel() == num_elem, "number of elements in v and p tensors should be equal");
+        AT_ASSERTM(g.numel() == num_elem, "number of elements in g and p tensors should be equal");
+        AT_ASSERTM(p_copy.numel() == num_elem || p_copy.numel() == 0, "number of elements in p_copy and p tensors should be equal, or p_copy should be empty");
+
+        fused_adam_cuda(p, p_copy, m, v, g, lr, beta1, beta2, eps, grad_scale, step, mode, bias_correction, decay);
+}
+void reversible_adam(at::Tensor & p, at::Tensor & p_copy, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay) {
+        CHECK_INPUT(p);
+        if (p_copy.numel() > 0) CHECK_INPUT(p_copy);
+        CHECK_INPUT(m);
+        CHECK_INPUT(v);
+        CHECK_INPUT(g);
+        int64_t num_elem = p.numel();
+        AT_ASSERTM(m.numel() == num_elem, "number of elements in m and p tensors should be equal");
+        AT_ASSERTM(v.numel() == num_elem, "number of elements in v and p tensors should be equal");
+        AT_ASSERTM(g.numel() == num_elem, "number of elements in g and p tensors should be equal");
+        AT_ASSERTM(p_copy.numel() == num_elem || p_copy.numel() == 0, "number of elements in p_copy and p tensors should be equal, or p_copy should be empty");
+
+        fused_reversible_adam_cuda(p, p_copy, m, v, g, lr, beta1, beta2, eps, grad_scale, step, mode, bias_correction, decay);
+}
+void maybe_adam_undo(at::Tensor & overflow_flag, at::Tensor & p, at::Tensor & m, at::Tensor & v, at::Tensor & g, float lr, float beta1, float beta2, float eps, float grad_scale, int step, int mode, int bias_correction, float decay) {
+        CHECK_INPUT(p);
+        CHECK_INPUT(m);
+        CHECK_INPUT(v);
+        CHECK_INPUT(g);
+        int64_t num_elem = p.numel();
+        AT_ASSERTM(m.numel() == num_elem, "number of elements in m and p tensors should be equal");
+        AT_ASSERTM(v.numel() == num_elem, "number of elements in v and p tensors should be equal");
+        AT_ASSERTM(g.numel() == num_elem, "number of elements in g and p tensors should be equal");
+
+        fused_maybe_adam_undo_cuda(overflow_flag, p, m, v, g, lr, beta1, beta2, eps, grad_scale, step, mode, bias_correction, decay);
+}
+void maybe_cast(at::Tensor & overflow_flag, at::Tensor & p_in, at::Tensor & p_out) {
+	CHECK_INPUT(p_in);
+	CHECK_INPUT(p_out);
+	int64_t num_elem = p_in.numel();
+	AT_ASSERTM(p_out.numel() == num_elem, "number of elements in p_in and p_out should be equal");
+
+	maybe_cast_cuda(overflow_flag, p_in, p_out);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+        m.def("strided_check_finite", &strided_check_finite, "Strided finite check.");
+        m.def("adam", &adam, "Adam optimized CUDA implementation.");
+        m.def("reversible_adam", &reversible_adam, "Reversible Adam optimized CUDA implementation.");
+        m.def("adam_mt", &fused_adam_cuda_mt, "Multi tensor Adam optimized CUDA implementation.");
+        m.def("maybe_adam_undo", &maybe_adam_undo, "Undo function for Adam optimized CUDA implementation.");
+        m.def("maybe_cast", &maybe_cast, "Unpack byte tensor containing e5m2 floats.");
+        m.def("maybe_cast_mt", &maybe_cast_cuda_mt, "Unpack byte tensor containing e5m2 floats.");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_adam_cuda_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_adam_cuda_kernel.cu
new file mode 100644
index 0000000000..75bd35a109
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_adam_cuda_kernel.cu
@@ -0,0 +1,1037 @@
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <stdio.h>
+#include <cmath>
+
+#include "ATen/ATen.h"
+#include "ATen/cuda/CUDAContext.h"
+#include "ATen/cuda/detail/IndexUtils.cuh"
+#include "ATen/TensorUtils.h"
+// #include "ATen/Type.h"
+#include "ATen/AccumulateType.h"
+
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+template<typename T>
+__device__ __forceinline__ bool is_aligned(T* p){
+  return ((uint64_t)p) % (ILP*sizeof(T)) == 0;
+}
+
+template<typename T>
+__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+
+#include "type_shim.h"
+
+typedef enum{
+    ADAM_MODE_0   =0, // eps under square root
+    ADAM_MODE_1   =1  // eps outside square root
+} adamMode_t;
+
+template <typename T, typename GRAD_T>
+__global__ void adam_cuda_kernel(
+        T* __restrict__ p,
+        GRAD_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed
+        T* __restrict__ m,
+        T* __restrict__ v,
+        const GRAD_T * __restrict__ g,
+        const float b1,
+        const float b2,
+        const float eps,
+        const float grad_scale,
+        const float step_size,
+        const size_t tsize,
+        adamMode_t mode,
+        const float decay)
+{
+        //Assuming 2D grids and 2D blocks
+        const int blockId = gridDim.x * blockIdx.y + blockIdx.x;
+        const int threadsPerBlock = blockDim.x * blockDim.y;
+        const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;
+        const int i = (blockId * threadsPerBlock + threadIdInBlock);
+        const int totThreads = gridDim.x*gridDim.y*threadsPerBlock;
+
+        for (int j = i; j < tsize; j+=totThreads) {
+                T scaled_grad = g[j]/grad_scale;
+                m[j] = b1*m[j] + (1-b1)*scaled_grad;
+                v[j] = b2*v[j] + (1-b2)*scaled_grad*scaled_grad;
+                float denom;
+                if (mode == ADAM_MODE_0)
+                    denom = sqrtf(v[j] + eps);
+                else // Mode 1
+                    denom = sqrtf(v[j]) + eps;
+                float update = (m[j]/denom) + (decay*p[j]);
+                p[j] = p[j] - (step_size*update);
+                if (p_copy != NULL) p_copy[j] = (GRAD_T) p[j];
+        }
+}
+
+template <int DEPTH, typename T, typename GRAD_T>
+struct AdamFunctor
+{
+    __device__ __forceinline__ void operator()(
+        int chunk_size,
+        volatile int* noop_gmem,
+        TensorListMetadata<DEPTH>& tl,
+        const float b1,
+        const float b2,
+        const float eps,
+        const float grad_scale,
+        const float step_size,
+        adamMode_t mode,
+        const float decay)
+    {
+        int tensor_loc = tl.block_to_tensor[blockIdx.x];
+        int chunk_idx = tl.block_to_chunk[blockIdx.x];
+        int n = tl.sizes[tensor_loc];
+
+        T* p = (T *)tl.addresses[0][tensor_loc];
+        p += chunk_idx*chunk_size;
+        T* m = (T *)tl.addresses[1][tensor_loc];
+        m += chunk_idx*chunk_size;
+        T* v = (T *)tl.addresses[2][tensor_loc];
+        v += chunk_idx*chunk_size;
+        GRAD_T* g = (GRAD_T *)tl.addresses[3][tensor_loc];
+        g += chunk_idx*chunk_size;
+        GRAD_T* p_copy = NULL;
+        if (DEPTH == 5) {
+            p_copy = (GRAD_T *)tl.addresses[4][tensor_loc];
+            p_copy += chunk_idx*chunk_size;
+        }
+
+        n -= chunk_idx*chunk_size;
+
+        T incoming_p[ILP];
+        T incoming_m[ILP];
+        T incoming_v[ILP];
+        T incoming_g[ILP];
+
+        // to make things simple, we put aligned case in a different code path
+        if(n % ILP == 0 &&
+           chunk_size % ILP == 0 &&
+           is_aligned(p) &&
+           is_aligned(m) &&
+           is_aligned(v) &&
+           is_aligned(g) &&
+           is_aligned(p_copy))
+        {
+          for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+          {
+            // load
+            GRAD_T tmp_g[ILP];
+            load_store(incoming_p, p, 0, i_start);
+            load_store(incoming_m, m, 0, i_start);
+            load_store(incoming_v, v, 0, i_start);
+            load_store(tmp_g, g, 0, i_start);
+#pragma unroll
+            for(int ii = 0; ii < ILP; ii++) {
+              incoming_g[ii] = static_cast<T>(tmp_g[ii]);
+              T scaled_grad = incoming_g[ii]/grad_scale;
+              incoming_m[ii] = b1*incoming_m[ii] + (1-b1)*scaled_grad;
+              incoming_v[ii] = b2*incoming_v[ii] + (1-b2)*scaled_grad*scaled_grad;
+              float denom;
+              if (mode == ADAM_MODE_0)
+                denom = sqrtf(incoming_v[ii] + eps);
+              else // Mode 1
+                denom = sqrtf(incoming_v[ii]) + eps;
+              float update = (incoming_m[ii]/denom) + (decay*incoming_p[ii]);
+              incoming_p[ii] = incoming_p[ii] - (step_size*update);
+              if (DEPTH == 5)  tmp_g[ii] = static_cast<GRAD_T>(incoming_p[ii]);
+            }
+            load_store(p, incoming_p, i_start, 0);
+            load_store(m, incoming_m, i_start, 0);
+            load_store(v, incoming_v, i_start, 0);
+            if (DEPTH == 5) load_store(p_copy, tmp_g, i_start, 0);
+          }
+        }
+        else
+        {
+          for(int i_start = 0;
+              i_start < n && i_start < chunk_size;
+              i_start += blockDim.x*ILP) {
+
+#pragma unroll
+            for(int ii = 0; ii < ILP; ii++) {
+              incoming_p[ii] = 0;
+              incoming_m[ii] = 0;
+              incoming_v[ii] = 0;
+              incoming_g[ii] = 0;
+
+              int i = i_start + threadIdx.x + ii*blockDim.x;
+              if (i < n && i < chunk_size) {
+                incoming_p[ii] = p[i];
+                incoming_m[ii] = m[i];
+                incoming_v[ii] = v[i];
+                incoming_g[ii] = static_cast<T>(g[i]);
+              }
+            }
+
+            // note for clarification to future michael:
+            // From a pure memory dependency perspective, there's likely no point unrolling
+            // the write loop, since writes just fire off once their LDGs arrive.
+            // Put another way, the STGs are dependent on the LDGs, but not on each other.
+            // There is still compute ILP benefit from unrolling the loop though.
+#pragma unroll
+            for(int ii = 0; ii < ILP; ii++) {
+              int j = i_start + threadIdx.x + ii*blockDim.x;
+
+              if(j < n && j < chunk_size) {
+                T scaled_grad = incoming_g[ii]/grad_scale;
+                m[j] = b1*incoming_m[ii] + (1-b1)*scaled_grad;
+                v[j] = b2*incoming_v[ii] + (1-b2)*scaled_grad*scaled_grad;
+                float denom;
+                if (mode == ADAM_MODE_0)
+                  denom = sqrtf(v[j] + eps);
+                else // Mode 1
+                  denom = sqrtf(v[j]) + eps;
+                float update = (m[j]/denom) + (decay*incoming_p[ii]);
+                p[j] = incoming_p[ii] - (step_size*update);
+                if (DEPTH == 5)  p_copy[j] = (GRAD_T) p[j];
+              }
+            }
+          }
+        }
+    }
+};
+
+void fused_adam_cuda(
+        at::Tensor & p,
+        at::Tensor & p_copy,
+        at::Tensor & m,
+        at::Tensor & v,
+        at::Tensor & g,
+        float lr,
+        float beta1,
+        float beta2,
+        float eps,
+        float grad_scale,
+        int step,
+        int mode,
+        int bias_correction,
+        float decay)
+{
+//        using namespace at;
+
+        //Get tensor size
+        int tsize = p.numel();
+        //Determine #threads and #blocks
+        const int threadsPerBlock = 512;
+        const dim3 blocks((tsize+threadsPerBlock-1)/threadsPerBlock);
+        AT_ASSERTM(at::cuda::detail::canUse32BitIndexMath(p), "parameter tensor is too large to be indexed with int32");
+        //Constants
+        float step_size = 0;
+        if (bias_correction == 1) {
+            const float bias_correction1 = 1 - std::pow(beta1, step);
+            const float bias_correction2 = 1 - std::pow(beta2, step);
+            step_size = lr * std::sqrt(bias_correction2)/bias_correction1;
+        }
+        else {
+            step_size = lr;
+        }
+        cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+        if (g.scalar_type() == at::ScalarType::Half) {
+//all other values should be fp32 for half gradients
+            AT_ASSERTM(p.scalar_type() == at::ScalarType::Float, "expected parameter to be of float type");
+//dispatch is done on the gradient type
+            using namespace at; // prevents "toString is undefined" errors
+            DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, "adam_cuda_kernel",
+                using accscalar_t = at::acc_type<scalar_t_0, true>;
+                adam_cuda_kernel<accscalar_t, scalar_t_0><<<blocks,threadsPerBlock, 0, stream>>>(
+                        p.DATA_PTR<accscalar_t>(),
+                        p_copy.numel() ? p_copy.DATA_PTR<scalar_t_0>() : NULL,
+                        m.DATA_PTR<accscalar_t>(),
+                        v.DATA_PTR<accscalar_t>(),
+                        g.DATA_PTR<scalar_t_0>(),
+                        beta1,
+                        beta2,
+                        eps,
+                        grad_scale,
+                        step_size,
+                        tsize,
+                        (adamMode_t) mode,
+                        decay);
+                );
+      } else {
+            using namespace at;
+            DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, "adam_cuda_kernel",
+                adam_cuda_kernel<scalar_t_0, scalar_t_0><<<blocks,threadsPerBlock, 0, stream>>>(
+                        p.DATA_PTR<scalar_t_0>(),
+                        NULL, //don't output p_copy for fp32, it's wasted write
+                        m.DATA_PTR<scalar_t_0>(),
+                        v.DATA_PTR<scalar_t_0>(),
+                        g.DATA_PTR<scalar_t_0>(),
+                        beta1,
+                        beta2,
+                        eps,
+                        grad_scale,
+                        step_size,
+                        tsize,
+                        (adamMode_t) mode,
+                        decay);
+            );
+      }
+      C10_CUDA_CHECK(cudaGetLastError());
+
+}
+
+void fused_adam_cuda_mt(
+    int chunk_size,
+    at::Tensor noop_flag,
+    std::vector<std::vector<at::Tensor>> tensor_lists, // p, m, v, g, p_copy
+    float lr,
+    float beta1,
+    float beta2,
+    float eps,
+    float grad_scale,
+    int step,
+    int mode,
+    int bias_correction,
+    float decay) {
+
+    //Constants
+    float step_size = 0;
+    if (bias_correction == 1) {
+        const float bias_correction1 = 1 - std::pow(beta1, step);
+        const float bias_correction2 = 1 - std::pow(beta2, step);
+        step_size = lr * std::sqrt(bias_correction2)/bias_correction1;
+    }
+    else {
+        step_size = lr;
+    }
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    size_t tl_sz = tensor_lists.size();
+    AT_ASSERTM(tl_sz == 4 || tl_sz == 5, "expected tensor lists of size 4 or 5");
+
+    if (tensor_lists[3][0].scalar_type() == at::ScalarType::Half) {
+//alher values should be fp32 for half gradients
+        AT_ASSERTM(tensor_lists[0][0].scalar_type() == at::ScalarType::Float, "expected parameter to be of float type");
+//dich is done on the gradient type
+        if (tl_sz == 5) {
+            DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, "adam_cuda_mt_kernel",
+                using accscalar_t = at::acc_type<scalar_t_0, true>;
+                multi_tensor_apply<5>(
+                    BLOCK_SIZE,
+                    chunk_size,
+                    noop_flag,
+                    tensor_lists,
+                    AdamFunctor<5, accscalar_t, scalar_t_0>(),
+                    beta1,
+                    beta2,
+                    eps,
+                    grad_scale,
+                    step_size,
+                    (adamMode_t) mode,
+                    decay);
+            );
+        } else {
+            DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, "adam_cuda_mt_kernel",
+                using accscalar_t = at::acc_type<scalar_t_0, true>;
+                multi_tensor_apply<4>(
+                    BLOCK_SIZE,
+                    chunk_size,
+                    noop_flag,
+                    tensor_lists,
+                    AdamFunctor<4, accscalar_t, scalar_t_0>(),
+                    beta1,
+                    beta2,
+                    eps,
+                    grad_scale,
+                    step_size,
+                    (adamMode_t) mode,
+                    decay);
+            );
+        }
+    } else {
+        if (tl_sz == 5) {
+            DISPATCH_DOUBLE_AND_FLOAT(tensor_lists[3][0].scalar_type(), 0, "adam_cuda_mt_kernel",
+                multi_tensor_apply<5>(
+                    BLOCK_SIZE,
+                    chunk_size,
+                    noop_flag,
+                    tensor_lists,
+                    AdamFunctor<5, scalar_t_0, scalar_t_0>(),
+                    beta1,
+                    beta2,
+                    eps,
+                    grad_scale,
+                    step_size,
+                    (adamMode_t) mode,
+                    decay);
+            );
+        } else {
+            DISPATCH_DOUBLE_AND_FLOAT(tensor_lists[3][0].scalar_type(), 0, "adam_cuda_mt_kernel",
+                multi_tensor_apply<4>(
+                    BLOCK_SIZE,
+                    chunk_size,
+                    noop_flag,
+                    tensor_lists,
+                    AdamFunctor<4, scalar_t_0, scalar_t_0>(),
+                    beta1,
+                    beta2,
+                    eps,
+                    grad_scale,
+                    step_size,
+                    (adamMode_t) mode,
+                    decay);
+            );
+        }
+    }
+    C10_CUDA_CHECK(cudaGetLastError());
+}
+
+template <typename FROM_T, typename TO_T> 
+__device__ void convert(const FROM_T vi, TO_T& vo)
+{
+    vo = static_cast<TO_T>(vi);
+}
+
+template <>
+__device__ void convert(const float vi, uint8_t& vo)
+{
+    union S
+    {
+	float as_float;
+	int as_int;
+    };
+    S s;
+    s.as_float = vi;
+    s.as_int = s.as_int & 0xFF800000;
+    union T
+    {
+        at::Half as_half;
+	uint8_t as_byte[2];
+    };
+    T t;
+    t.as_half = static_cast<at::Half>(vi + s.as_float / 8.0f);
+    vo = t.as_byte[1];
+}
+
+template <>
+__device__ void convert(const uint8_t vi, float& vo)
+{
+    union T
+    {
+        at::Half as_half;
+	uint8_t as_byte[2];
+    };
+    T t;
+    t.as_byte[0] = 0;
+    t.as_byte[1] = vi;
+    vo = static_cast<float>(t.as_half);
+}
+
+template <>
+__device__ void convert(const at::Half vi, uint8_t& vo)
+{
+    union S
+    {
+	float as_float;
+	int as_int;
+    };
+    S s;
+    s.as_float = static_cast<float>(vi);
+    s.as_int = s.as_int & 0xFF800000;
+    union T
+    {
+        at::Half as_half;
+	uint8_t as_byte[2];
+    };
+    T t;
+    t.as_half = static_cast<at::Half>(vi + s.as_float / 8.0f);
+    vo = t.as_byte[1];
+}
+
+template <>
+__device__ void convert(const uint8_t vi, at::Half& vo)
+{
+    union T
+    {
+        at::Half as_half;
+	uint8_t as_byte[2];
+    };
+    T t;
+    t.as_byte[0] = 0;
+    t.as_byte[1] = vi;
+    vo = t.as_half;
+}
+
+template <typename GRAD_T>
+__global__ void strided_check_finite_cuda_kernel(
+        volatile int* noop_gmem,
+        GRAD_T* __restrict__ p_copy,
+        const size_t tsize,
+        int stride,
+        int clear_overflow_first)
+{
+    //Assuming 2D grids and 2D blocks
+    const int blockId = gridDim.x * blockIdx.y + blockIdx.x;
+    const int threadsPerBlock = blockDim.x * blockDim.y;
+    const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;
+    const int i = (blockId * threadsPerBlock + threadIdInBlock) * stride;
+    const int totThreads = gridDim.x*gridDim.y*threadsPerBlock*stride;
+
+    if (clear_overflow_first) {
+        if (i == 0) {
+            *noop_gmem = 0;
+        }
+        __syncthreads();
+    }
+
+    for (int j = i; j < tsize; j+=totThreads) {
+        GRAD_T pi = p_copy[j];
+        if (!isfinite(pi)) {
+            *noop_gmem = 1;
+        }
+    }
+}
+template <>
+__global__ void strided_check_finite_cuda_kernel(
+        volatile int* noop_gmem,
+        uint8_t* __restrict__ p_copy,
+        const size_t tsize,
+        int stride,
+        int clear_overflow_first)
+{
+    //Assuming 2D grids and 2D blocks
+    const int blockId = gridDim.x * blockIdx.y + blockIdx.x;
+    const int threadsPerBlock = blockDim.x * blockDim.y;
+    const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;
+    const int i = (blockId * threadsPerBlock + threadIdInBlock) * stride;
+    const int totThreads = gridDim.x*gridDim.y*threadsPerBlock*stride;
+
+    if (clear_overflow_first) {
+        if (i == 0) {
+            *noop_gmem = 0;
+        }
+        __syncthreads();
+    }
+
+    for (int j = i; j < tsize; j+=totThreads) {
+        at::Half pi;
+        convert(p_copy[j], pi);
+        if (!isfinite(pi)) {
+            *noop_gmem = 1;
+        }
+    }
+}
+
+template <typename FROM_T, typename TO_T> 
+__global__ void maybe_cast_kernel(
+        volatile int* overflow_flag,
+        const FROM_T* p_in,
+        TO_T* p_out,
+        const size_t tsize)
+{
+    if (overflow_flag && *overflow_flag != 0) return;
+
+    //Assuming 2D grids and 2D blocks
+    const int blockId = gridDim.x * blockIdx.y + blockIdx.x;
+    const int threadsPerBlock = blockDim.x * blockDim.y;
+    const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;
+    const int i = (blockId * threadsPerBlock + threadIdInBlock);
+    const int totThreads = gridDim.x*gridDim.y*threadsPerBlock;
+
+    FROM_T pi[ILP];
+    TO_T po[ILP];
+
+    for(int j_start = 0;  j_start < tsize;  j_start+=totThreads*ILP) {
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++) {
+            pi[ii] = 0;
+
+            int j = j_start + i + totThreads*ii;
+            if (j < tsize) {
+                pi[ii] = p_in[j];
+            }
+        }
+
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++) {
+            convert(pi[ii], po[ii]);
+        }
+
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++) {
+            int j = j_start + i + totThreads*ii;
+            if (j < tsize) {
+                p_out[j] = po[ii];
+            }
+        }
+    }
+}
+
+template <typename T, typename GRAD_T, typename REDU_T>
+__global__ void reversible_adam_cuda_kernel(
+        T* __restrict__ p,
+        REDU_T* __restrict__ p_copy, // For mixed precision training, pass NULL if not needed
+        T* __restrict__ m,
+        T* __restrict__ v,
+        const GRAD_T * __restrict__ g,
+        const float b1,
+        const float b2,
+        const float eps,
+        const float grad_scale,
+        const float step_size,
+        const size_t tsize,
+        adamMode_t mode,
+        const float decay)
+{
+    //Assuming 2D grids and 2D blocks
+    const int blockId = gridDim.x * blockIdx.y + blockIdx.x;
+    const int threadsPerBlock = blockDim.x * blockDim.y;
+    const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;
+    const int i = (blockId * threadsPerBlock + threadIdInBlock);
+    const int totThreads = gridDim.x*gridDim.y*threadsPerBlock;
+
+    T mi[ILP];
+    T vi[ILP];
+    T pi[ILP];
+    T gi[ILP];
+
+    bool overflow = false;
+    for(int j_start = 0;  j_start < tsize;  j_start+=totThreads*ILP) {
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++) {
+            mi[ii] = T(0);
+            vi[ii] = T(0);
+            pi[ii] = T(0);
+            gi[ii] = GRAD_T(0);
+
+            int j = j_start + i + totThreads*ii;
+            if (j < tsize) {
+                pi[ii] = p[j];
+                mi[ii] = m[j];
+                vi[ii] = v[j];
+                gi[ii] = static_cast<T>(g[j]);
+            }
+        }
+
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++) {
+            T scaled_grad = gi[ii]/grad_scale;
+            if (isfinite(scaled_grad)) {
+                mi[ii] = b1*mi[ii] + (1-b1)*scaled_grad;
+                vi[ii] = b2*vi[ii] + (1-b2)*scaled_grad*scaled_grad;
+                float denom;
+                if (mode == ADAM_MODE_0)
+                    denom = sqrtf(vi[ii] + eps);
+                else // Mode 1
+                    denom = sqrtf(vi[ii]) + eps;
+                float update = (mi[ii]/denom) + (decay*pi[ii]);
+                pi[ii] = pi[ii] - (step_size*update);
+            } else {
+                overflow = true;
+            }
+        }
+
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++) {
+            int j = j_start + i + totThreads*ii;
+            if (j < tsize) {
+                m[j] = mi[ii];
+                v[j] = vi[ii];
+                p[j] = pi[ii];
+                if (p_copy != NULL) {
+                    convert(pi[ii], p_copy[j]);
+                }
+            }
+        }
+    }
+
+    if (p_copy != NULL) {
+        __syncthreads();
+        if (overflow) {
+            convert(float(INFINITY), p_copy[0]);
+        }
+    }
+}
+
+template <typename T, typename GRAD_T>
+__global__ void maybe_adam_undo_cuda_kernel(
+        volatile int* overflow_flag,
+        T* __restrict__ p,
+        T* __restrict__ m,
+        T* __restrict__ v,
+        const GRAD_T * __restrict__ g,
+        const float b1,
+        const float b2,
+        const float eps,
+        const float grad_scale,
+        const float step_size,
+        const size_t tsize,
+        adamMode_t mode,
+        const float decay)
+{
+    // NB! Skip undo kernel when overflow flag is NOT set
+    if (overflow_flag && *overflow_flag == 0) return;
+
+    //Assuming 2D grids and 2D blocks
+    const int blockId = gridDim.x * blockIdx.y + blockIdx.x;
+    const int threadsPerBlock = blockDim.x * blockDim.y;
+    const int threadIdInBlock = threadIdx.y * blockDim.x + threadIdx.x;
+    const int i = (blockId * threadsPerBlock + threadIdInBlock);
+    const int totThreads = gridDim.x*gridDim.y*threadsPerBlock;
+
+    T mi[ILP];
+    T vi[ILP];
+    T pi[ILP];
+    T gi[ILP];
+
+    for(int j_start = 0;  j_start < tsize;  j_start+=totThreads*ILP) {
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++) {
+            mi[ii] = T(0);
+            vi[ii] = T(0);
+            pi[ii] = T(0);
+            gi[ii] = GRAD_T(0);
+
+            int j = j_start + i*ILP;
+            if (j < tsize) {
+                pi[ii] = p[j];
+                mi[ii] = m[j];
+                vi[ii] = v[j];
+                gi[ii] = static_cast<T>(g[j]);
+            }
+        }
+
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++) {
+            T scaled_grad = gi[ii]/grad_scale;
+            if (isfinite(scaled_grad)) {
+                float denom;
+                if (mode == ADAM_MODE_0)
+                    denom = sqrtf(vi[ii] + eps);
+                else // Mode 1
+                    denom = sqrtf(vi[ii]) + eps;
+                pi[ii] = (pi[ii] + step_size*(mi[ii]/denom)) / (1.0f - step_size*decay);
+                mi[ii] = (mi[ii] - (1-b1)*scaled_grad) / b1;
+                vi[ii] = (vi[ii] - (1-b2)*scaled_grad*scaled_grad) / b2;
+                // Make sure round off errors don't create (small) negative value.
+                // This can happen if we have to revert the very first step.
+                vi[ii] = vi[ii] >= 0.0f ? vi[ii] : 0.0f;
+            }
+        }
+
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++) {
+            int j = j_start + i*ILP;
+            if (j < tsize) {
+                m[j] = mi[ii];
+                v[j] = vi[ii];
+                p[j] = pi[ii];
+            }
+        }
+    }
+}
+
+template <int DEPTH, typename FROM_T, typename TO_T>
+struct MaybeCastFunctor
+{
+    __device__ __forceinline__ void operator()(
+        int chunk_size,
+        volatile int* overflow_flag,
+        TensorListMetadata<DEPTH>& tl)
+    {
+        if (overflow_flag && *overflow_flag != 0) return;
+
+        int tensor_loc = tl.block_to_tensor[blockIdx.x];
+        int chunk_idx = tl.block_to_chunk[blockIdx.x];
+        int n = tl.sizes[tensor_loc];
+
+        FROM_T* p_in = (FROM_T *)tl.addresses[0][tensor_loc];
+        p_in += chunk_idx*chunk_size;
+        TO_T* p_out = (TO_T *)tl.addresses[1][tensor_loc];
+        p_out += chunk_idx*chunk_size;
+
+        n -= chunk_idx*chunk_size;
+        int dim = chunk_size < n ? chunk_size : n;
+
+	FROM_T pi[ILP];
+        TO_T po[ILP];
+
+        for(int j_start = 0;  j_start < dim;  j_start+=blockDim.x*ILP) {
+#pragma unroll
+            for(int ii = 0; ii < ILP; ii++) {
+                pi[ii] = FROM_T(0);
+                int j = j_start + threadIdx.x + ii*blockDim.x;
+                if (j < dim) {
+                    pi[ii] = p_in[j];
+                }
+            }
+
+#pragma unroll
+            for(int ii = 0; ii < ILP; ii++) {
+                convert(pi[ii], po[ii]);
+            }
+
+#pragma unroll
+            for(int ii = 0; ii < ILP; ii++) {
+                int j = j_start + threadIdx.x + ii*blockDim.x;
+                if (j < dim) {
+                    p_out[j] = po[ii];
+                }
+            }
+        }
+    }
+};
+
+void fused_strided_check_finite(
+	at::Tensor & overflow_flag,
+        at::Tensor & p_copy,
+        int stride,
+	int clear_overflow_first)
+{
+	//Get tensor size
+	int tsize = p_copy.numel();
+	int niter = (tsize + stride - 1) / stride;
+
+	//Determine #threads and #blocks
+	const int threadsPerBlock = 512;
+	//In order to avoid race condition, blocks must be 1 when clear_overflow_first flag is set.
+	const dim3 blocks(clear_overflow_first ? 1 : (niter+threadsPerBlock-1)/threadsPerBlock);
+	AT_ASSERTM(at::cuda::detail::canUse32BitIndexMath(p_copy), "parameter tensor is too large to be indexed with int32");
+
+	cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+        using namespace at; // prevents "toString is undefined" errors
+        DISPATCH_FLOAT_HALF_AND_BYTE(p_copy.scalar_type(), 0, "check_finite_cuda_kernel",
+                strided_check_finite_cuda_kernel<scalar_t_0><<<blocks,threadsPerBlock, 0, stream>>>(
+                    overflow_flag.DATA_PTR<int>(),
+                    p_copy.DATA_PTR<scalar_t_0>(),
+                    tsize,
+                    stride,
+                    clear_overflow_first);
+                );
+	C10_CUDA_CHECK(cudaGetLastError());
+}
+
+void fused_reversible_adam_cuda(
+        at::Tensor & p,
+        at::Tensor & p_copy,
+        at::Tensor & m,
+        at::Tensor & v,
+        at::Tensor & g,
+        float lr,
+        float beta1,
+        float beta2,
+        float eps,
+        float grad_scale,
+        int step,
+        int mode,
+        int bias_correction,
+        float decay)
+{
+//      using namespace at;
+
+      //Get tensor size
+      int tsize = p.numel();
+      //Determine #threads and #blocks
+      const int threadsPerBlock = 512;
+      const dim3 blocks((tsize+threadsPerBlock-1)/threadsPerBlock);
+      AT_ASSERTM(at::cuda::detail::canUse32BitIndexMath(p), "parameter tensor is too large to be indexed with int32");
+      //Constants
+      float step_size = 0;
+      if (bias_correction == 1) {
+          const float bias_correction1 = 1 - std::pow(beta1, step);
+          const float bias_correction2 = 1 - std::pow(beta2, step);
+          step_size = lr * std::sqrt(bias_correction2)/bias_correction1;
+      }
+      else {
+          step_size = lr;
+      }
+      cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+      if (g.scalar_type() == at::ScalarType::Half) {
+          //all other values should be fp32 for half gradients
+          AT_ASSERTM(p.scalar_type() == at::ScalarType::Float, "expected parameter to be of float type");
+          //dispatch is done on the gradient type
+          using namespace at; // prevents "toString is undefined" errors
+          if (p_copy.numel() == 0 || p_copy.scalar_type() == g.scalar_type()) {
+              DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, "adam_cuda_kernel",
+                      using accscalar_t = at::acc_type<scalar_t_0, true>;
+                      reversible_adam_cuda_kernel<accscalar_t, scalar_t_0, scalar_t_0><<<blocks,threadsPerBlock, 0, stream>>>(
+                          p.DATA_PTR<accscalar_t>(),
+                          p_copy.numel() ? p_copy.DATA_PTR<scalar_t_0>() : NULL,
+                          m.DATA_PTR<accscalar_t>(),
+                          v.DATA_PTR<accscalar_t>(),
+                          g.DATA_PTR<scalar_t_0>(),
+                          beta1,
+                          beta2,
+                          eps,
+                          grad_scale,
+                          step_size,
+                          tsize,
+                          (adamMode_t) mode,
+                          decay);
+                      );
+          } else {
+              AT_ASSERTM(p_copy.scalar_type() == at::ScalarType::Byte, "expected parameter to be of byte type");
+              DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, "adam_cuda_e5m2_kernel",
+                      using accscalar_t = at::acc_type<scalar_t_0, true>;
+                      reversible_adam_cuda_kernel<accscalar_t, scalar_t_0, uint8_t><<<blocks,threadsPerBlock, 0, stream>>>(
+                          p.DATA_PTR<accscalar_t>(),
+                          p_copy.DATA_PTR<uint8_t>(),
+                          m.DATA_PTR<accscalar_t>(),
+                          v.DATA_PTR<accscalar_t>(),
+                          g.DATA_PTR<scalar_t_0>(),
+                          beta1,
+                          beta2,
+                          eps,
+                          grad_scale,
+                          step_size,
+                          tsize,
+                          (adamMode_t) mode,
+                          decay);
+                      );
+          }
+      } else {
+          using namespace at;
+          DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, "adam_cuda_kernel",
+                  reversible_adam_cuda_kernel<scalar_t_0, scalar_t_0, scalar_t_0><<<blocks,threadsPerBlock, 0, stream>>>(
+                      p.DATA_PTR<scalar_t_0>(),
+                      NULL, //don't output p_copy for fp32, it's wasted write
+                      m.DATA_PTR<scalar_t_0>(),
+                      v.DATA_PTR<scalar_t_0>(),
+                      g.DATA_PTR<scalar_t_0>(),
+                      beta1,
+                      beta2,
+                      eps,
+                      grad_scale,
+                      step_size,
+                      tsize,
+                      (adamMode_t) mode,
+                      decay);
+                  );
+      }
+      C10_CUDA_CHECK(cudaGetLastError());
+}
+
+void maybe_cast_cuda(
+        at::Tensor & overflow_flag,
+        at::Tensor & p_in,
+        at::Tensor & p_out)
+{
+      //Get tensor size
+      int tsize = p_in.numel();
+      AT_ASSERTM(tsize == p_out.numel(), "p_in.numel() must equal p_out.numel()");
+      //Determine #threads and #blocks
+      const int threadsPerBlock = 512;
+      const dim3 blocks((tsize+threadsPerBlock-1)/threadsPerBlock);
+      AT_ASSERTM(at::cuda::detail::canUse32BitIndexMath(p_in), "parameter tensor is too large to be indexed with int32");
+      //Constants
+      cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+      DISPATCH_FLOAT_HALF_AND_BYTE(p_in.scalar_type(), 0, "maybe_cast_cuda"
+              DISPATCH_FLOAT_HALF_AND_BYTE(p_out.scalar_type(), 1, "maybe_cast_cuda",
+                  maybe_cast_kernel<scalar_t_0,scalar_t_1><<<blocks,threadsPerBlock, 0, stream>>>(
+                      overflow_flag.numel() ? overflow_flag.DATA_PTR<int>() : NULL,
+                      p_in.DATA_PTR<scalar_t_0>(),
+                      p_out.DATA_PTR<scalar_t_1>(),
+                      tsize); ))
+      C10_CUDA_CHECK(cudaGetLastError());
+}
+
+void maybe_cast_cuda_mt(
+    int chunk_size,
+    at::Tensor overflow_flag,
+    std::vector<std::vector<at::Tensor>> tensor_lists) // p_in, p_out
+{
+    //Constants
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    size_t tl_sz = tensor_lists.size();
+    AT_ASSERTM(tl_sz == 2, "expected tensor lists of size 2");
+
+    DISPATCH_FLOAT_HALF_AND_BYTE(tensor_lists[0][0].scalar_type(), 0, "maybe_cast_cuda_mt_kernel",
+            DISPATCH_FLOAT_HALF_AND_BYTE(tensor_lists[1][0].scalar_type(), 1, "maybe_cast_cuda_mt_kernel",
+                multi_tensor_apply<2>(
+                    BLOCK_SIZE,
+                    chunk_size,
+                    overflow_flag,
+                    tensor_lists,
+                    MaybeCastFunctor<2, scalar_t_0, scalar_t_1>()); ))
+    C10_CUDA_CHECK(cudaGetLastError());
+}
+
+void fused_maybe_adam_undo_cuda(
+        at::Tensor & overflow_flag,
+        at::Tensor & p,
+        at::Tensor & m,
+        at::Tensor & v,
+        at::Tensor & g,
+        float lr,
+        float beta1,
+        float beta2,
+        float eps,
+        float grad_scale,
+        int step,
+        int mode,
+        int bias_correction,
+        float decay)
+{
+    //Get tensor size
+    int tsize = p.numel();
+    //Determine #threads and #blocks
+    const int threadsPerBlock = 512;
+    const dim3 blocks((tsize+threadsPerBlock-1)/threadsPerBlock);
+    AT_ASSERTM(at::cuda::detail::canUse32BitIndexMath(p), "parameter tensor is too large to be indexed with int32");
+    //Constants
+    float step_size = 0;
+    if (bias_correction == 1) {
+        const float bias_correction1 = 1 - std::pow(beta1, step);
+        const float bias_correction2 = 1 - std::pow(beta2, step);
+        step_size = lr * std::sqrt(bias_correction2)/bias_correction1;
+    }
+    else {
+        step_size = lr;
+    }
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    if (g.scalar_type() == at::ScalarType::Half) {
+        //all other values should be fp32 for half gradients
+        AT_ASSERTM(p.scalar_type() == at::ScalarType::Float, "expected parameter to be of float type");
+        //dispatch is done on the gradient type
+        using namespace at; // prevents "toString is undefined" errors
+        DISPATCH_FLOAT_AND_HALF(g.scalar_type(), 0, "adam_cuda_kernel",
+                using accscalar_t = at::acc_type<scalar_t_0, true>;
+                maybe_adam_undo_cuda_kernel<accscalar_t, scalar_t_0><<<blocks,threadsPerBlock, 0, stream>>>(
+                    overflow_flag.numel() ? overflow_flag.DATA_PTR<int>() : NULL,
+                    p.DATA_PTR<accscalar_t>(),
+                    m.DATA_PTR<accscalar_t>(),
+                    v.DATA_PTR<accscalar_t>(),
+                    g.DATA_PTR<scalar_t_0>(),
+                    beta1,
+                    beta2,
+                    eps,
+                    grad_scale,
+                    step_size,
+                    tsize,
+                    (adamMode_t) mode,
+                    decay);
+                );
+    } else {
+        using namespace at;
+        DISPATCH_DOUBLE_AND_FLOAT(g.scalar_type(), 0, "adam_cuda_kernel",
+                maybe_adam_undo_cuda_kernel<scalar_t_0, scalar_t_0><<<blocks,threadsPerBlock, 0, stream>>>(
+                    overflow_flag.numel() ? overflow_flag.DATA_PTR<int>() : NULL,
+                    p.DATA_PTR<scalar_t_0>(),
+                    m.DATA_PTR<scalar_t_0>(),
+                    v.DATA_PTR<scalar_t_0>(),
+                    g.DATA_PTR<scalar_t_0>(),
+                    beta1,
+                    beta2,
+                    eps,
+                    grad_scale,
+                    step_size,
+                    tsize,
+                    (adamMode_t) mode,
+                    decay);
+                );
+    }
+    C10_CUDA_CHECK(cudaGetLastError());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_lamb_cuda.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_lamb_cuda.cpp
new file mode 100644
index 0000000000..98a24115ca
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_lamb_cuda.cpp
@@ -0,0 +1,21 @@
+#include <torch/extension.h>
+
+void multi_tensor_lamb_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  const float lr,
+  const float beta1,
+  const float beta2,
+  const float epsilon,
+  const int step,
+  const int bias_correction,
+  const float weight_decay,
+  const int grad_averaging,
+  const int mode,
+  const float global_grad_norm,
+  const float max_grad_norm);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+        m.def("lamb", &multi_tensor_lamb_cuda, "Computes and apply update for LAMB optimizer");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_lamb_cuda_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_lamb_cuda_kernel.cu
new file mode 100644
index 0000000000..3bb93b0312
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/fused_lamb_cuda_kernel.cu
@@ -0,0 +1,294 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+typedef enum{
+  MOMENT_MODE_0   =0, // L2 regularization mode
+  MOMENT_MODE_1   =1  // Decoupled weight decay mode
+} adamMode_t;
+
+std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::optional<bool> per_tensor_python);
+
+using MATH_T = float;
+
+template<typename T>
+struct LAMBStage1Functor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<4>& tl,
+    const float beta1,
+    const float beta2,
+    const float beta3,
+    const float beta1_correction,
+    const float beta2_correction,
+    const float epsilon,
+    adamMode_t mode,
+    const float decay,
+    const float global_grad_norm,
+    const float max_global_grad_norm)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    float clipped_global_grad_norm = global_grad_norm > max_global_grad_norm ? global_grad_norm / max_global_grad_norm : 1.0f;
+
+    T* g = (T*)tl.addresses[0][tensor_loc];
+    g += chunk_idx*chunk_size;
+
+    T* p = (T*)tl.addresses[1][tensor_loc];
+    p += chunk_idx*chunk_size;
+
+    T* m = (T*)tl.addresses[2][tensor_loc];
+    m += chunk_idx*chunk_size;
+
+    T* v = (T*)tl.addresses[3][tensor_loc];
+    v += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    // see note in multi_tensor_scale_kernel.cu
+    for(int i_start = 0;
+            i_start < n && i_start < chunk_size;
+            i_start += blockDim.x*ILP)
+    {
+      MATH_T r_g[ILP];
+      MATH_T r_p[ILP];
+      MATH_T r_m[ILP];
+      MATH_T r_v[ILP];
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          r_g[ii] = g[i];
+          // special ?optimization? for lamb stage 1
+          if (decay == 0) {
+            r_p[ii] = MATH_T(0);
+          }
+          else {
+            r_p[ii] = p[i];
+          }
+          r_m[ii] = m[i];
+          r_v[ii] = v[i];
+        } else {
+          r_g[ii] = MATH_T(0);
+          r_p[ii] = MATH_T(0);
+          r_m[ii] = MATH_T(0);
+          r_v[ii] = MATH_T(0);
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        if (mode == MOMENT_MODE_0) {
+	  MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
+	  // L2 on scaled grad
+          scaled_grad = scaled_grad + decay*r_p[ii];
+          r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
+          r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+          MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+          MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+          MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+          r_p[ii] = next_m_unbiased / denom;
+        }
+        else {
+          MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
+          r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
+          r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+          MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+          MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+          MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+          r_p[ii] = (next_m_unbiased/denom) + (decay*r_p[ii]);
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          g[i] = r_p[ii];
+          m[i] = r_m[ii];
+          v[i] = r_v[ii];
+        }
+      }
+    }
+  }
+};
+
+// Step 2 reads in 'update' value and per-tensor param_norm and update_norm.
+// It computes new parameter value.
+template<typename T>
+struct LAMBStage2Functor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<2>& tl,
+    const float* per_tensor_param_norm,
+    const float* per_tensor_update_norm,
+    const float learning_rate,
+    const float decay)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int tensor_num = tl.start_tensor_this_launch + tensor_loc;
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    MATH_T ratio = learning_rate;
+    // apply adaptive learning rate to parameters with non-zero weight decay
+    if (decay != 0.0) 
+    {
+      float param_norm = per_tensor_param_norm[tensor_num];
+      float update_norm = per_tensor_update_norm[tensor_num];
+      ratio = (update_norm != 0.0f && param_norm != 0.0f) ? learning_rate * (param_norm / update_norm) : learning_rate;
+    }
+
+    T* update = (T*)tl.addresses[0][tensor_loc];
+    update += chunk_idx*chunk_size;
+
+    T* p = (T*)tl.addresses[1][tensor_loc];
+    p += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    for(int i_start = 0;
+            i_start < n && i_start < chunk_size;
+            i_start += blockDim.x*ILP)
+    {
+      MATH_T r_p[ILP];
+      MATH_T r_update[ILP];
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+       	int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          r_p[ii] = p[i];
+          r_update[ii] = update[i];
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+       	r_p[ii] = r_p[ii] - (ratio * r_update[ii]);
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          p[i] = r_p[ii];
+        }
+      }
+    }
+  }
+};
+
+
+void multi_tensor_lamb_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  const float lr,
+  const float beta1,
+  const float beta2,
+  const float epsilon,
+  const int step,
+  const int bias_correction,
+  const float weight_decay,
+  const int grad_averaging,
+  const int mode,
+  const float global_grad_norm,
+  const float max_grad_norm)
+{
+  using namespace at;
+  // Master weight and 32bit momentum(potentially changing) is not handled by this
+  // So we assume every tensor are all in the same type
+
+  // Handle bias correction mode
+  float bias_correction1 = 1.0f, bias_correction2 = 1.0f;
+  if (bias_correction == 1) {
+    bias_correction1 = 1 - std::pow(beta1, step);
+    bias_correction2 = 1 - std::pow(beta2, step);
+  }
+
+  // Handle grad averaging mode
+  float beta3 = 1.0f;
+  if (grad_averaging == 1) beta3 = 1 - beta1;
+
+  std::vector<std::vector<at::Tensor>> grad_list(tensor_lists.begin(), tensor_lists.begin()+1);
+  std::vector<std::vector<at::Tensor>> param_list(tensor_lists.begin()+1, tensor_lists.begin()+2);
+
+  // Compute per tensor param norm
+  auto param_norm_tuple = multi_tensor_l2norm_cuda(chunk_size, noop_flag, param_list, true);
+
+  // We now in-place modify grad to store update before compute its norm
+  // Generally this is not a issue since people modify grad in step() method all the time
+  // We can also grab list of empty tensor to avoid this, but I'd like to save space/cpu code
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "lamb_stage_1",
+      multi_tensor_apply<4>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        LAMBStage1Functor<scalar_t_0>(),
+        beta1,
+        beta2,
+        beta3, // 1-beta1 or 1 depends on averaging mode
+        bias_correction1,
+        bias_correction2,
+        epsilon,
+        (adamMode_t) mode,
+        weight_decay,
+        global_grad_norm,
+        max_grad_norm); )
+
+  // Compute update norms
+  auto update_norm_tuple = multi_tensor_l2norm_cuda(chunk_size, noop_flag, grad_list, true);
+
+  std::vector<std::vector<at::Tensor>> grad_param_list(tensor_lists.begin(), tensor_lists.begin()+2);
+
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "lamb_stage_2",
+      multi_tensor_apply<2>(
+        BLOCK_SIZE,
+        chunk_size,
+       	noop_flag,
+        grad_param_list,
+        LAMBStage2Functor<scalar_t_0>(),
+        std::get<1>(param_norm_tuple).DATA_PTR<float>(),
+        std::get<1>(update_norm_tuple).DATA_PTR<float>(),
+        lr,
+	weight_decay); )
+
+  AT_CUDA_CHECK(cudaGetLastError());
+
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_adam.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_adam.cpp
new file mode 100644
index 0000000000..7ae13d5149
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_adam.cpp
@@ -0,0 +1,20 @@
+#include <torch/extension.h>
+
+void multi_tensor_fused_adam_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor per_tensor_beta1,
+  at::Tensor per_tensor_beta2,
+  at::Tensor per_tensor_bias_correction,
+  at::Tensor per_tensor_eps,
+  at::Tensor per_tensor_weight_decay,
+  float lr,
+  float grad_scale,
+  int step,
+  int mode);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("multi_tensor_fused_adam", &multi_tensor_fused_adam_cuda,
+        "Multi tensor Adam optimized CUDA implementation.");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_adam_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_adam_kernel.cu
new file mode 100644
index 0000000000..a702adab66
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_adam_kernel.cu
@@ -0,0 +1,228 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+#include <cmath>
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+template<typename T>
+__device__ __forceinline__ bool is_aligned(T* p){
+  return ((uint64_t)p) % (ILP*sizeof(T)) == 0;
+}
+
+template<typename T>
+__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+
+typedef enum{
+  ADAM_MODE_0   =0, // eps under square root
+  ADAM_MODE_1   =1  // eps outside square root
+} adamMode_t;
+
+template <int DEPTH, typename T, typename GRAD_T>
+struct DistAdamFunctor
+{
+  __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<DEPTH>& tl,
+    const float* per_tensor_beta1,
+    const float* per_tensor_beta2,
+    const int* per_tensor_bias_correction,
+    const float* per_tensor_eps,
+    const float* per_tensor_weight_decay,
+    const float lr,
+    const float grad_scale,
+    const int step,
+    adamMode_t mode)
+  {
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int tensor_num = tl.start_tensor_this_launch + tensor_loc;
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    float b1 = per_tensor_beta1[tensor_num];
+    float b2 = per_tensor_beta2[tensor_num];
+    float eps = per_tensor_eps[tensor_num];
+    float decay = per_tensor_weight_decay[tensor_num];
+
+    float beta1_correction = 1.0f, beta2_correction = 1.0f;
+    if (per_tensor_bias_correction[tensor_num] == 1) {
+      beta1_correction = 1 - std::pow(b1, step);
+      beta2_correction = 1 - std::pow(b2, step);
+    }
+
+    T* p = (T *)tl.addresses[0][tensor_loc];
+    p += chunk_idx*chunk_size;
+    T* m = (T *)tl.addresses[1][tensor_loc];
+    m += chunk_idx*chunk_size;
+    T* v = (T *)tl.addresses[2][tensor_loc];
+    v += chunk_idx*chunk_size;
+    GRAD_T* g = (GRAD_T *)tl.addresses[3][tensor_loc];
+    g += chunk_idx*chunk_size;
+    GRAD_T* p_copy = NULL;
+    if (DEPTH == 5) {
+      p_copy = (GRAD_T *)tl.addresses[4][tensor_loc];
+      p_copy += chunk_idx*chunk_size;
+    }
+
+    n -= chunk_idx*chunk_size;
+    
+    T incoming_p[ILP];
+    T incoming_m[ILP];
+    T incoming_v[ILP];
+    T incoming_g[ILP];
+
+    // to make things simple, we put aligned case in a different code path
+    if (n % ILP == 0 &&
+      chunk_size % ILP == 0 &&
+      is_aligned(p) &&
+      is_aligned(m) &&
+      is_aligned(v) &&
+      is_aligned(g) &&
+      is_aligned(p_copy)) {
+      for (int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x) {
+        // load
+        GRAD_T tmp_g[ILP];
+        load_store(incoming_p, p, 0, i_start);
+        load_store(incoming_m, m, 0, i_start);
+        load_store(incoming_v, v, 0, i_start);
+        load_store(tmp_g, g, 0, i_start);
+#pragma unroll
+        for (int ii = 0; ii < ILP; ii++) {
+          incoming_g[ii] = static_cast<T>(tmp_g[ii]);
+          T scaled_grad = incoming_g[ii]/grad_scale;
+          incoming_m[ii] = b1*incoming_m[ii] + (1-b1)*scaled_grad;
+          incoming_v[ii] = b2*incoming_v[ii] + (1-b2)*scaled_grad*scaled_grad;
+          T next_m_unbiased = incoming_m[ii] / beta1_correction;
+	  T next_v_unbiased = incoming_v[ii] / beta2_correction;
+	  float denom;
+          if (mode == ADAM_MODE_0)
+            denom = sqrtf(next_v_unbiased + eps);
+          else // Mode 1
+            denom = sqrtf(next_v_unbiased) + eps;
+          float update = (next_m_unbiased / denom) + (decay * incoming_p[ii]);
+          incoming_p[ii] = incoming_p[ii] - (lr * update);
+	  if (DEPTH == 5)  tmp_g[ii] = static_cast<GRAD_T>(incoming_p[ii]);
+        }
+        load_store(p, incoming_p, i_start, 0);
+        load_store(m, incoming_m, i_start, 0);
+        load_store(v, incoming_v, i_start, 0);
+        if (DEPTH == 5) load_store(p_copy, tmp_g, i_start, 0);
+      }
+    } else {
+      for (int i_start = 0;
+          i_start < n && i_start < chunk_size;
+          i_start += blockDim.x*ILP) {
+
+#pragma unroll
+        for (int ii = 0; ii < ILP; ii++) {
+          incoming_p[ii] = 0;
+          incoming_m[ii] = 0;
+          incoming_v[ii] = 0;
+          incoming_g[ii] = 0;
+
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if (i < n && i < chunk_size) {
+            incoming_p[ii] = p[i];
+            incoming_m[ii] = m[i];
+            incoming_v[ii] = v[i];
+            incoming_g[ii] = static_cast<T>(g[i]);
+          }
+        }
+
+#pragma unroll
+        for (int ii = 0; ii < ILP; ii++) {
+          int j = i_start + threadIdx.x + ii*blockDim.x;
+
+          if (j < n && j < chunk_size) {
+            T scaled_grad = incoming_g[ii]/grad_scale;
+            m[j] = b1*incoming_m[ii] + (1-b1)*scaled_grad;
+            v[j] = b2*incoming_v[ii] + (1-b2)*scaled_grad*scaled_grad;
+            T next_m_unbiased = m[j] / beta1_correction;
+            T next_v_unbiased = v[j] / beta2_correction;
+	    float denom;
+            if (mode == ADAM_MODE_0)
+              denom = sqrtf(next_v_unbiased + eps);
+            else // Mode 1
+              denom = sqrtf(next_v_unbiased) + eps;
+            float update = (next_m_unbiased / denom) + (decay * incoming_p[ii]);
+            p[j] = incoming_p[ii] - (lr * update);
+	    if (DEPTH == 5)  p_copy[j] = (GRAD_T) p[j];
+          }
+        }
+      }
+    }
+  }
+};
+
+void multi_tensor_fused_adam_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,  // p, m, v, g, p_copy
+  at::Tensor per_tensor_beta1,
+  at::Tensor per_tensor_beta2,
+  at::Tensor per_tensor_bias_correction,
+  at::Tensor per_tensor_eps,
+  at::Tensor per_tensor_weight_decay,
+  float lr,
+  float grad_scale,
+  int step,
+  int mode)
+{
+  using namespace at;
+
+  size_t tl_sz = tensor_lists.size();
+  AT_ASSERTM(tl_sz == 4 || tl_sz == 5, "expected tensor lists of size 4 or 5");
+
+  if (tl_sz == 5) {
+    DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, "dist_adam_cuda_kernel",  // g
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      multi_tensor_apply<5>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        DistAdamFunctor<5, accscalar_t, scalar_t_0>(),
+        per_tensor_beta1.DATA_PTR<float>(),
+        per_tensor_beta2.DATA_PTR<float>(),
+        per_tensor_bias_correction.DATA_PTR<int>(),
+        per_tensor_eps.DATA_PTR<float>(),
+        per_tensor_weight_decay.DATA_PTR<float>(),
+        lr,
+        grad_scale,
+        step,
+        (adamMode_t) mode);
+    );
+  } else {
+    DISPATCH_FLOAT_AND_HALF(tensor_lists[3][0].scalar_type(), 0, "dist_adam_cuda_kernel",  // g
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      multi_tensor_apply<4>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        DistAdamFunctor<4, accscalar_t, scalar_t_0>(),
+        per_tensor_beta1.DATA_PTR<float>(),
+        per_tensor_beta2.DATA_PTR<float>(),
+        per_tensor_bias_correction.DATA_PTR<int>(),
+        per_tensor_eps.DATA_PTR<float>(),
+        per_tensor_weight_decay.DATA_PTR<float>(),
+        lr,
+        grad_scale,
+        step,
+        (adamMode_t) mode);
+    );
+  }
+  C10_CUDA_CHECK(cudaGetLastError());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb.cpp
new file mode 100644
index 0000000000..584b2a0e75
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb.cpp
@@ -0,0 +1,36 @@
+#include <torch/extension.h>
+
+void multi_tensor_lamb_compute_update_term_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor per_tensor_beta1,
+  at::Tensor per_tensor_beta2,
+  at::Tensor per_tensor_beta3,
+  at::Tensor per_tensor_bias_correction,
+  at::Tensor step,
+  at::Tensor per_tensor_epsilon,
+  const int mode,
+  at::Tensor per_tensor_decay,
+  at::Tensor global_scale,
+  at::Tensor global_grad_norm,
+  const float max_grad_norm);
+
+void multi_tensor_lamb_update_weights_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor per_tensor_param_norm,
+  at::Tensor per_tensor_update_norm,
+  at::Tensor update_norm_offset,
+  at::Tensor learning_rate,
+  at::Tensor per_tensor_decay,
+  at::Tensor global_grad_norm,
+  bool use_nvlamb);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("multi_tensor_lamb_compute_update_term", &multi_tensor_lamb_compute_update_term_cuda,
+        "Computes update term for LAMB optimizer");
+  m.def("multi_tensor_lamb_update_weights", &multi_tensor_lamb_update_weights_cuda,
+        "Applies update term for LAMB optimizer");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb_kernel.cu
new file mode 100644
index 0000000000..95ee009b2b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb_kernel.cu
@@ -0,0 +1,506 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+template<typename T>
+__device__ __forceinline__ bool is_aligned(T* p){
+  return ((uint64_t)p) % (ILP*sizeof(T)) == 0;
+}
+
+template<typename T>
+__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+
+template <typename FROM_T, typename TO_T> 
+__device__ void convert(const FROM_T vi, TO_T& vo)
+{
+    vo = static_cast<TO_T>(vi);
+}
+
+template <>
+__device__ void convert(const float vi, uint8_t& vo)
+{
+    union S
+    {
+	float as_float;
+	int as_int;
+    };
+    S s;
+    s.as_float = vi;
+    s.as_int = s.as_int & 0xFF800000;
+    union T
+    {
+        at::Half as_half;
+	uint8_t as_byte[2];
+    };
+    T t;
+    t.as_half = static_cast<at::Half>(vi + s.as_float / 8.0f);
+    vo = t.as_byte[1];
+}
+
+template <>
+__device__ void convert(const uint8_t vi, float& vo)
+{
+    union T
+    {
+        at::Half as_half;
+	uint8_t as_byte[2];
+    };
+    T t;
+    t.as_byte[0] = 0;
+    t.as_byte[1] = vi;
+    vo = static_cast<float>(t.as_half);
+}
+
+template <>
+__device__ void convert(const at::Half vi, uint8_t& vo)
+{
+    union S
+    {
+	float as_float;
+	int as_int;
+    };
+    S s;
+    s.as_float = static_cast<float>(vi);
+    s.as_int = s.as_int & 0xFF800000;
+    union T
+    {
+        at::Half as_half;
+	uint8_t as_byte[2];
+    };
+    T t;
+    t.as_half = static_cast<at::Half>(vi + s.as_float / 8.0f);
+    vo = t.as_byte[1];
+}
+
+template <>
+__device__ void convert(const uint8_t vi, at::Half& vo)
+{
+    union T
+    {
+        at::Half as_half;
+	uint8_t as_byte[2];
+    };
+    T t;
+    t.as_byte[0] = 0;
+    t.as_byte[1] = vi;
+    vo = t.as_half;
+}
+
+typedef enum{
+  MOMENT_MODE_0   =0, // L2 regularization mode
+  MOMENT_MODE_1   =1  // Decoupled weight decay mode
+} adamMode_t;
+
+template<typename T, typename GRAD_T, typename MATH_T>
+struct DistOptLAMBStage1Functor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<5>& tl,
+    const MATH_T* per_tensor_beta1,
+    const MATH_T* per_tensor_beta2,
+    const MATH_T* per_tensor_beta3,
+    const int* per_tensor_bias_correction,
+    const int* step,
+    const MATH_T* per_tensor_epsilon,
+    adamMode_t mode,
+    const MATH_T* per_tensor_decay,
+    const MATH_T* global_scale,
+    const MATH_T* global_grad_norm,
+    const float max_grad_norm)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    if (*noop_gmem == 1)
+        return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int tensor_num = tl.start_tensor_this_launch + tensor_loc;
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    float combined_scale = *global_scale;
+    if (max_grad_norm > 0) {
+        combined_scale = max_grad_norm / (*global_grad_norm / *global_scale + 1e-6);
+	combined_scale = *global_scale / std::min((float) 1.0, combined_scale);
+    }
+
+    MATH_T beta1 = per_tensor_beta1[tensor_num];
+    MATH_T beta2 = per_tensor_beta2[tensor_num];
+    MATH_T beta3 = 1 - beta1;
+    MATH_T beta1_correction, beta2_correction;
+    if (per_tensor_bias_correction[tensor_num] == 1) {
+        beta1_correction = 1 - pow(beta1, *step);
+        beta2_correction = 1 - pow(beta2, *step);
+    } else {
+        beta1_correction = (MATH_T) 1.0;
+        beta2_correction = (MATH_T) 1.0;
+    }
+    MATH_T epsilon = per_tensor_epsilon[tensor_num];
+    MATH_T decay = per_tensor_decay[tensor_num];
+
+    GRAD_T* g = (GRAD_T*)tl.addresses[0][tensor_loc];
+    g += chunk_idx*chunk_size;
+
+    T* p = (T*)tl.addresses[1][tensor_loc];
+    p += chunk_idx*chunk_size;
+
+    T* m = (T*)tl.addresses[2][tensor_loc];
+    m += chunk_idx*chunk_size;
+
+    T* v = (T*)tl.addresses[3][tensor_loc];
+    v += chunk_idx*chunk_size;
+
+    MATH_T* u = (MATH_T*)tl.addresses[4][tensor_loc];
+    u += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    MATH_T r_g[ILP];
+    MATH_T r_p[ILP];
+    MATH_T r_m[ILP];
+    MATH_T r_v[ILP];
+    // to make things simple, we put aligned case in a different code path
+    if(n % ILP == 0 &&
+       chunk_size % ILP == 0 &&
+       is_aligned(g) &&
+       is_aligned(p) &&
+       is_aligned(m) &&
+       is_aligned(v))
+    {
+      GRAD_T l_g[ILP];
+      T l_p[ILP];
+      T l_m[ILP];
+      T l_v[ILP];
+      for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+      {
+        // load
+        load_store(l_g, g, 0, i_start);
+        if (decay != 0)
+          load_store(l_p, p, 0, i_start);
+        load_store(l_m, m, 0, i_start);
+        load_store(l_v, v, 0, i_start);
+        // unpack
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_g[ii] = l_g[ii];
+          if (decay == 0) {
+            r_p[ii] = MATH_T(0);
+          }
+          else {
+            r_p[ii] = l_p[ii];
+          }
+          r_m[ii] = l_m[ii];
+          r_v[ii] = l_v[ii];
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          if (mode == MOMENT_MODE_0) {
+            MATH_T scaled_grad = r_g[ii] / combined_scale;
+            // L2 on scaled grad
+            scaled_grad = scaled_grad + decay*r_p[ii];
+            r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
+            r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+            MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+            MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+            MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+            r_p[ii] = next_m_unbiased / denom;
+          }
+          else {
+            MATH_T scaled_grad = r_g[ii] / combined_scale;
+            r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
+            r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+            MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+            MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+            MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+            r_p[ii] = (next_m_unbiased/denom) + (decay*r_p[ii]);
+          }
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          l_m[ii] = r_m[ii];
+          l_v[ii] = r_v[ii];
+        }
+        // store
+        load_store(u, r_p, i_start, 0);
+        load_store(m, l_m, i_start, 0);
+        load_store(v, l_v, i_start, 0);
+      }
+    }
+    else
+    {
+      // see note in multi_tensor_scale_kernel.cu
+      for(int i_start = 0;
+          i_start < n && i_start < chunk_size;
+          i_start += blockDim.x*ILP)
+      {
+        MATH_T r_g[ILP];
+        MATH_T r_p[ILP];
+        MATH_T r_m[ILP];
+        MATH_T r_v[ILP];
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            r_g[ii] = g[i];
+            // special ?optimization? for lamb stage 1
+            if (decay == 0) {
+              r_p[ii] = MATH_T(0);
+            }
+            else {
+              r_p[ii] = p[i];
+            }
+            r_m[ii] = m[i];
+            r_v[ii] = v[i];
+          } else {
+            r_g[ii] = MATH_T(0);
+            r_p[ii] = MATH_T(0);
+            r_m[ii] = MATH_T(0);
+            r_v[ii] = MATH_T(0);
+          }
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          if (mode == MOMENT_MODE_0) {
+            MATH_T scaled_grad = r_g[ii] / combined_scale;
+            // L2 on scaled grad
+            scaled_grad = scaled_grad + decay*r_p[ii];
+            r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
+            r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+            MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+            MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+            MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+            r_p[ii] = next_m_unbiased / denom;
+          }
+          else {
+            MATH_T scaled_grad = r_g[ii] / combined_scale;
+            r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
+            r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+            MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+            MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+            MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+            r_p[ii] = (next_m_unbiased/denom) + (decay*r_p[ii]);
+          }
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            u[i] = r_p[ii];
+            m[i] = r_m[ii];
+            v[i] = r_v[ii];
+          }
+        }
+      }
+    }
+  }
+};
+
+// Step 2 reads in 'update' value and per-tensor param_norm and update_norm.
+// It computes new parameter value.
+template<typename T, typename GRAD_T, typename MATH_T>
+struct DistOptLAMBStage2Functor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<3>& tl,
+    const MATH_T* per_tensor_param_norm,
+    const MATH_T* per_tensor_update_norm,
+    const long* update_norm_offset,
+    const MATH_T* learning_rate,
+    const MATH_T* per_tensor_decay,
+    const MATH_T* global_grad_norm,
+    bool use_nvlamb)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    if (*noop_gmem == 1)
+        return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int tensor_num = tl.start_tensor_this_launch + tensor_loc;
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    MATH_T decay = per_tensor_decay[tensor_num];
+
+    MATH_T ratio = *learning_rate;
+    // nvlamb: apply adaptive learning rate to all parameters
+    // otherwise, only apply to those with non-zero weight decay
+    if (use_nvlamb || (decay != (MATH_T) 0.0))
+    {
+      MATH_T param_norm = per_tensor_param_norm[tensor_num];
+      MATH_T update_norm = per_tensor_update_norm[update_norm_offset[tensor_num]];
+      ratio = (update_norm != 0.0 && param_norm != 0.0) ? (*learning_rate) * (param_norm / update_norm) : (*learning_rate);
+    }
+
+    MATH_T* update = (MATH_T*)tl.addresses[0][tensor_loc];
+    update += chunk_idx*chunk_size;
+
+    T* p = (T*)tl.addresses[1][tensor_loc];
+    p += chunk_idx*chunk_size;
+
+    GRAD_T* p_copy = (GRAD_T*)tl.addresses[2][tensor_loc];
+    p_copy += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    // to make things simple, we put aligned case in a different code path
+    if(n % ILP == 0 &&
+       chunk_size % ILP == 0 &&
+       is_aligned(p) &&
+       is_aligned(update))
+    {
+      T r_p[ILP];
+      MATH_T r_update[ILP];
+      GRAD_T r_p_copy[ILP];
+      for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+      {
+        // load
+        load_store(r_p, p, 0, i_start);
+        load_store(r_update, update, 0, i_start);
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+	  r_p[ii] = static_cast<MATH_T>(r_p[ii]) - (ratio * r_update[ii]);
+          convert(r_p[ii], r_p_copy[ii]);
+        }
+        load_store(p, r_p, i_start, 0);
+        load_store(p_copy, r_p_copy, i_start, 0);
+      }
+    }
+    else
+    {
+      for(int i_start = 0;
+          i_start < n && i_start < chunk_size;
+          i_start += blockDim.x*ILP)
+      {
+        MATH_T r_p[ILP];
+        MATH_T r_update[ILP];
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            r_p[ii] = p[i];
+            r_update[ii] = update[i];
+          }
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_p[ii] = r_p[ii] - (ratio * r_update[ii]);
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            p[i] = r_p[ii];
+            convert(r_p[ii], p_copy[i]);
+          }
+        }
+      }
+    }
+  }
+};
+
+void multi_tensor_lamb_compute_update_term_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor per_tensor_beta1,
+  at::Tensor per_tensor_beta2,
+  at::Tensor per_tensor_beta3,
+  at::Tensor per_tensor_bias_correction,
+  at::Tensor step,
+  at::Tensor per_tensor_epsilon,
+  const int mode,
+  at::Tensor per_tensor_decay,
+  at::Tensor global_scale,
+  at::Tensor global_grad_norm,
+  const float max_grad_norm)
+{
+  using namespace at;
+
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[1][0].scalar_type(), 0, "lamb_stage_1",
+    DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 1, "lamb_stage_1",
+      DISPATCH_FLOAT_AND_HALF(tensor_lists[4][0].scalar_type(), 2, "lamb_stage_1",
+        multi_tensor_apply<5>(
+          BLOCK_SIZE,
+          chunk_size,
+          noop_flag,
+          tensor_lists,
+          DistOptLAMBStage1Functor<scalar_t_0, scalar_t_1, scalar_t_2>(),
+          per_tensor_beta1.DATA_PTR<scalar_t_2>(),
+          per_tensor_beta2.DATA_PTR<scalar_t_2>(),
+          per_tensor_beta3.DATA_PTR<scalar_t_2>(),
+          per_tensor_bias_correction.DATA_PTR<int>(),
+          step.DATA_PTR<int>(),
+          per_tensor_epsilon.DATA_PTR<scalar_t_2>(),
+          (adamMode_t) mode,
+          per_tensor_decay.DATA_PTR<scalar_t_2>(),
+          global_scale.DATA_PTR<scalar_t_2>(),
+	  global_grad_norm.DATA_PTR<scalar_t_2>(),
+	  max_grad_norm); )))
+
+  AT_CUDA_CHECK(cudaGetLastError());
+}
+
+void multi_tensor_lamb_update_weights_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor per_tensor_param_norm,
+  at::Tensor per_tensor_update_norm,
+  at::Tensor update_norm_offset,
+  at::Tensor learning_rate,
+  at::Tensor per_tensor_decay,
+  at::Tensor global_grad_norm,
+  bool use_nvlamb)
+{
+  using namespace at;
+
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[1][0].scalar_type(), 0, "lamb_stage_2",
+    DISPATCH_FLOAT_HALF_AND_BYTE(tensor_lists[2][0].scalar_type(), 1, "lamb_stage_2",
+      DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 2, "lamb_stage_2",
+        multi_tensor_apply<3>(
+          BLOCK_SIZE,
+          chunk_size,
+          noop_flag,
+          tensor_lists,
+          DistOptLAMBStage2Functor<scalar_t_0, scalar_t_1, scalar_t_2>(),
+          per_tensor_param_norm.DATA_PTR<scalar_t_2>(),
+          per_tensor_update_norm.DATA_PTR<scalar_t_2>(),
+          update_norm_offset.DATA_PTR<long>(),
+	  learning_rate.DATA_PTR<scalar_t_2>(),
+          per_tensor_decay.DATA_PTR<scalar_t_2>(),
+	  global_grad_norm.DATA_PTR<scalar_t_2>(),
+          use_nvlamb); )))
+
+  AT_CUDA_CHECK(cudaGetLastError());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_joint.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_joint.cpp
new file mode 100644
index 0000000000..ddbcd08062
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_joint.cpp
@@ -0,0 +1,98 @@
+#include <torch/extension.h>
+#include <ATen/Functions.h>
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> transducer_joint_cuda_forward(
+    torch::Tensor f,
+    torch::Tensor g,
+    torch::Tensor fLen,
+    torch::Tensor gLen,
+    torch::Tensor batchOffset,
+    int64_t packedBatch,
+    int opt,
+    bool packOutput,
+    bool relu,
+    bool dropout,
+    float dropoutProb,
+    int tileSize);
+
+
+std::vector<torch::Tensor> transducer_joint_cuda_backward(
+    std::vector<torch::Tensor> in,
+    torch::Tensor fLen,
+    torch::Tensor gLen,
+    torch::Tensor batchOffset,
+    int maxFLen,
+    int maxGLen,
+    bool packOutput,
+    float scale);
+
+std::vector<torch::Tensor> transducer_joint_forward(
+    torch::Tensor f,
+    torch::Tensor g,
+    torch::Tensor fLen,
+    torch::Tensor gLen,
+    torch::Tensor batchOffset,
+    int64_t packedBatch,
+    int opt,
+    bool packOutput,
+    bool relu,
+    bool dropout,
+    float dropoutProb,
+    int tileSize) {
+    CHECK_INPUT(f);
+    CHECK_INPUT(g);
+    CHECK_INPUT(fLen);
+    CHECK_INPUT(gLen);
+    if (packOutput)
+        CHECK_INPUT(batchOffset);
+    return transducer_joint_cuda_forward(
+        f, 
+        g, 
+        fLen, 
+        gLen,
+        batchOffset,
+        packedBatch,
+        opt,
+        packOutput,
+        relu,
+        dropout,
+        dropoutProb,
+        tileSize);
+}
+
+std::vector<torch::Tensor> transducer_joint_backward(
+    std::vector<torch::Tensor> in,
+    torch::Tensor fLen,
+    torch::Tensor gLen,
+    torch::Tensor batchOffset,
+    int maxFLen,
+    int maxGLen,
+    bool packOutput,
+    float scale) {
+    for (auto t : in){
+        CHECK_INPUT(t);
+    }
+    CHECK_INPUT(fLen);
+    CHECK_INPUT(gLen);
+    if (packOutput)
+        CHECK_INPUT(batchOffset);
+    return transducer_joint_cuda_backward(
+        in, 
+        fLen, 
+        gLen,
+        batchOffset,
+        maxFLen,
+        maxGLen,
+        packOutput,
+        scale);
+}
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &transducer_joint_forward, "transducer joint forward (CUDA)");
+  m.def("backward", &transducer_joint_backward, "transducer joint backward (CUDA)");
+}
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_joint_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_joint_kernel.cu
new file mode 100644
index 0000000000..2439a7f58a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_joint_kernel.cu
@@ -0,0 +1,973 @@
+#include <cuda.h>
+#include <curand_kernel.h>
+#include <cuda_runtime.h>
+
+#include <torch/extension.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/CUDAGeneratorImpl.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/CUDAGraphsUtils.cuh>
+#include <c10/macros/Macros.h>
+
+#include "philox.h"
+
+// Warp reduce kernels to reduce N groups of data into N numbers, where N = warpSize / width.
+// width should be a power of 2 and should be less than warpSize.
+template <typename scalar_t>
+__device__ __forceinline__ scalar_t warpReduce(scalar_t x, int width=C10_WARP_SIZE){
+    for (unsigned offset = width/2; offset > 0; offset /= 2){
+        x += __shfl_down_sync(0xffffffff, x, offset, width);   
+    }
+    return x;
+}
+
+inline int largestPowerOfTwo(int x){
+    int y = 1;
+    while (y <= x)
+        y <<= 1;
+    return y >> 1;
+}
+
+/*
+Figure out vectorization type for masks.
+Similar to how PyTorch figures out acc_t here:
+aten/src/ATen/AccumulateType.h 
+*/
+template <int V>
+struct MaskVecType { };
+
+template <> struct MaskVecType<1> { using type = uint8_t; };
+template <> struct MaskVecType<2> { using type = uint16_t; };
+template <> struct MaskVecType<4> { using type = uint32_t; };
+
+template<int V>
+using mvec_type = typename MaskVecType<V>::type;
+
+// Helper class to calculate pointer offset that can be shared by different flavors of kernels.
+// For fwd, batch offset and stride are different for packing and non-packing mode.
+struct OffsetCalFwd{
+    __device__ __forceinline__ OffsetCalFwd(
+        int64_t batch, 
+        const int64_t *batchOffset, 
+        int64_t maxFLen, 
+        int64_t maxGLen, 
+        int64_t gLen,
+        int64_t hiddenSize,
+        bool packOutput) :
+        batch(batch),
+        batchOffset(batchOffset),
+        maxFLen(maxFLen),
+        maxGLen(maxGLen),
+        gLen(gLen),
+        hiddenSize(hiddenSize),
+        packOutput(packOutput)
+        {}
+    
+    int64_t batch;
+    const int64_t *batchOffset;
+    int64_t maxFLen;
+    int64_t maxGLen;
+    int64_t gLen;
+    int64_t hiddenSize;
+    bool packOutput;
+
+    __device__ __forceinline__ int64_t getBatchOffset(){
+        return packOutput ? ((batch==0) ? 0 : batchOffset[batch-1])*hiddenSize 
+                            : batch*maxFLen*maxGLen*hiddenSize;
+    }
+
+    __device__ __forceinline__ int64_t getStrideF(){
+        return packOutput ? gLen*hiddenSize : maxGLen*hiddenSize;
+    }
+
+    
+};
+
+// Helper class to calculate pointer offset that can be shared by different flavors of kernels
+// For bwd, batch offset and stride are different for packing and non-packing mode.
+// The reducion is done for two input tensors. Therefore, generating two sets of offsets
+// according to bwdFasterDim can lead to a unified implementation in the actual kernel.
+struct OffsetCalBwd{
+    __device__ __forceinline__ OffsetCalBwd(
+        int64_t batch, 
+        const int64_t *batchOffset, 
+        const int *fLen, 
+        const int *gLen,
+        int64_t maxFLen, 
+        int64_t maxGLen, 
+        int64_t hiddenSize,
+        bool packOutput,
+        bool bwdFasterDim) :
+        batch(batch),
+        batchOffset(batchOffset),
+        maxFLen(maxFLen),
+        maxGLen(maxGLen),
+        fLen(fLen),
+        gLen(gLen),
+        hiddenSize(hiddenSize),
+        packOutput(packOutput),
+        bwdFasterDim(bwdFasterDim)
+        {}
+
+    int64_t batch;
+    const int64_t *batchOffset;
+    const int *fLen;
+    const int *gLen;
+    int64_t maxFLen;
+    int64_t maxGLen;
+    int64_t hiddenSize;
+    bool packOutput;
+    bool bwdFasterDim;  // whether doing bwd on the faster moving dimension
+
+    __device__ __forceinline__ int64_t getBatchOffset(){
+        return packOutput ? ((batch==0) ? 0 : batchOffset[batch-1])*hiddenSize 
+                            : batch*maxFLen*maxGLen*hiddenSize;
+    }
+
+    __device__ __forceinline__ int64_t getMaxXLen(){
+        return bwdFasterDim ? maxGLen : maxFLen;
+    }
+
+    __device__ __forceinline__ auto getMyXLen() -> decltype(gLen[batch]){
+        return bwdFasterDim ? gLen[batch] : fLen[batch];
+    }
+
+    __device__ __forceinline__ auto getMyYLen() -> decltype(gLen[batch]){
+        return bwdFasterDim ? fLen[batch] : gLen[batch];
+    }
+    
+    __device__ __forceinline__ int64_t getStrideX(){
+        return bwdFasterDim ? hiddenSize : ((packOutput ? gLen[batch] : maxGLen) * hiddenSize);
+    }
+
+    __device__ __forceinline__ int64_t getStrideY(){
+        return bwdFasterDim ? ((packOutput ? gLen[batch] : maxGLen) * hiddenSize) : hiddenSize;
+    }
+};
+
+
+// Vanila transducer joint forward kernel
+// Detail of this joint function can be found in: 
+// [1] Sequence Transduction with Recurrent Neural Networks.
+
+// f is a tensor of shape [batch, T, H]
+// g is a tensor of shape [batch, U, H]
+// the transducer joint does
+// sum = f.unsqueeze(dim=2) + g.unsqueeze(dim=1)
+// The resultant tensor is of shape [batch, T, U, H]
+// Each thread block is working on one "batch" of data in the output tensor, [batch, t, u, :]
+
+// This joint function can optionally pack the output where the output tensor with a shape of
+// [B, T, U, H] is packed into [B_packed, H].
+// Don't-care region (t > fLen) or (u > gLen) is removed.
+// To enable packing, the starting offset for each batch need to be specified with batchOffset.
+template <typename scalar_t, class OffsetCal>
+__global__ void transducer_joint_forward(
+    const scalar_t *f,
+    const scalar_t *g,
+    const int *fLen,
+    const int *gLen,
+    const int64_t *batchOffset,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    int64_t hiddenSize,
+    bool packOutput,
+    scalar_t *sum) {
+
+
+    const int batch = blockIdx.z;
+    const int t = blockIdx.y;
+    const int u = blockIdx.x;
+    const auto myFLen = fLen[batch];
+    const auto myGLen = gLen[batch];
+
+    OffsetCal offsetCal(batch, batchOffset, maxFLen, maxGLen, myGLen, hiddenSize, packOutput);
+    const auto myBatchOffset = offsetCal.getBatchOffset();
+    const auto strideF = offsetCal.getStrideF();
+    scalar_t const *myF = f + batch*maxFLen*hiddenSize + t*hiddenSize;
+    scalar_t const *myG = g + batch*maxGLen*hiddenSize + u*hiddenSize;
+    scalar_t *mySum = sum + myBatchOffset + t*strideF + u * hiddenSize;
+
+    if (t < myFLen and u < myGLen){
+        #pragma unroll
+        for (int h = threadIdx.x; h < hiddenSize; h += blockDim.x){
+            if (h < hiddenSize){
+                mySum[h] = myF[h] + myG[h];
+            }
+        }
+    }
+    else if (packOutput == false and t < maxFLen and u < maxGLen){
+        // Need to write finite data to don't-care region because we instantiate the result tensor
+        // with torch::empty for performance reasons. Even though it is don't-care region, the 
+        // contents need to be finite, otherwise could lead to NaN in WGRAD.
+        // In packing mode, this write is no longer necessary as we remove the don't-care region
+        // from the output.
+        // Picking -1 (over 0) here for ease of testing.
+        #pragma unroll
+        for (int h = threadIdx.x; h < hiddenSize; h += blockDim.x){
+            if (h < hiddenSize){
+                mySum[h] = -1;
+            }
+        }    
+    }
+}
+
+/*
+Tiled version of the joint forward kernel
+Detail of this joint function can be found in: 
+[1] Sequence Transduction with Recurrent Neural Networks.
+
+f is a tensor of shape [batch, T, H]
+g is a tensor of shape [batch, U, H]
+the transducer joint does
+sum = f.unsqueeze(dim=2) + g.unsqueeze(dim=1)
+The resultant tensor is of shape [batch, T, U, H]
+Each thread is working on a tile of the shape of tileF x tileG in the result tensor.
+The input for the tile is first loaded in the register and is reused tileG and tileF times. 
+
+This joint function can optionally pack the output where the output tensor with a shape of
+[B, T, U, H] is packed into [B_packed, H].
+Don't-care region (t > fLen) or (u > gLen) is removed.
+To enable packing, the starting offset for each batch need to be specified with batchOffset.
+
+Optionally this joint function performs ReLU and/or dropout on the joint output, which is 
+controlled by arguments relu and dropout, respectively. philoxArgs is argument used for generating
+pseudorandom number. When at least one of operations in ReLU and dropout is activated, the joint
+function is a masked operation, which is controlled by the template argument masked. In this case, 
+masks are saved to backward.
+*/
+template <typename scalar_t, int tileF, int tileG, int U, class OffsetCal, bool masked>
+__global__ void transducer_joint_tiled_forward(
+    const scalar_t *f,
+    const scalar_t *g,
+    const int *fLen,
+    const int *gLen,
+    const int64_t *batchOffset,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    int64_t hiddenSize,
+    int64_t hiddenPerBlock,
+    bool packOutput,
+    bool relu, 
+    bool dropout,
+    float p,
+    at::PhiloxCudaState philoxArgs,
+    scalar_t *sum,
+    uint8_t *mask) {
+
+    static_assert(U == 4, "U has to be 4, as random numbers are generated in batch of 4");
+
+    const int batch = blockIdx.z;
+    const int t = blockIdx.y * tileF;
+    const int hiddenBlock = (hiddenSize + hiddenPerBlock - 1) / hiddenPerBlock;
+    const int u = blockIdx.x / hiddenBlock * tileG;
+    const int hOffset = (blockIdx.x % hiddenBlock) * hiddenPerBlock;
+    const int h = threadIdx.x;
+    const auto myFLen = fLen[batch];
+    const auto myGLen = gLen[batch];
+
+    OffsetCal offsetCal(batch, batchOffset, maxFLen, maxGLen, myGLen, hiddenSize, packOutput);
+    const auto myBatchOffset = offsetCal.getBatchOffset();
+    const auto strideF = offsetCal.getStrideF();
+
+    scalar_t const *myF = f + batch*maxFLen*hiddenSize + t*hiddenSize + hOffset;
+    scalar_t const *myG = g + batch*maxGLen*hiddenSize + u*hiddenSize + hOffset;
+    scalar_t *mySum = sum + myBatchOffset + t*strideF + u*hiddenSize + hOffset;
+    uint8_t *myMask = mask + myBatchOffset + t*strideF + u*hiddenSize + hOffset;
+
+    // The following code is only needed for dropout. We try to bypass them as much as possible.
+    auto seeds = masked ? at::cuda::philox::unpack(philoxArgs) 
+                            : std::make_tuple(static_cast<uint64_t>(0), static_cast<uint64_t>(0));
+    uint64_t tid = masked ? (static_cast<uint64_t>(blockIdx.z)*gridDim.y*gridDim.x + 
+                        blockIdx.y*gridDim.x + blockIdx.x) * blockDim.x + threadIdx.x
+                            : 0;
+    Philox ph(std::get<0>(seeds), tid, std::get<1>(seeds)); 
+    scalar_t scale = masked ? ((p == 0) ? 0 : 1 / p) : 0;  
+    bool dropoutMask[U];
+
+    if (t < myFLen and u < myGLen and hOffset+h < hiddenSize){    
+        // register buffers for tiled input reuse
+        scalar_t fBuffer[tileF], gBuffer[tileG];    
+        for (int i = 0; i < tileF; ++i){
+            if (t + i < myFLen)
+                fBuffer[i] = myF[i*hiddenSize + h];
+        }
+        for (int j = 0; j < tileG; ++j){
+            if (u + j < myGLen)
+                gBuffer[j] = myG[j*hiddenSize + h];
+        }
+        #pragma unroll
+        for (int i = 0; i < tileF; ++i){
+            if (t + i < myFLen){
+                #pragma unroll
+                for (int j = 0; j < tileG; ++j){
+                    int idx = i*tileG + j;
+                    if (masked and dropout and idx % U == 0){
+                        // For performance, generate 4 random numbers in one shot
+                        // auto rand4 = curand_uniform4(&state);
+                        auto rand4 = uniform4(ph());
+                        dropoutMask[0] = rand4.x < p;
+                        dropoutMask[1] = rand4.y < p;
+                        dropoutMask[2] = rand4.z < p;
+                        dropoutMask[3] = rand4.w < p;
+                    }
+
+                    if (u + j < myGLen){
+                        scalar_t out = fBuffer[i] + gBuffer[j];
+                        if (masked){
+                            // Apply ReLU here when relu is True
+                            bool localMask = relu ? (out>0) : 1;
+                            localMask = dropout ? localMask & dropoutMask[idx%U] : localMask;
+                            out = dropout ? out*localMask*scale : out*localMask;
+                            myMask[i*strideF + j*hiddenSize + h] = static_cast<uint8_t>(localMask);
+                        }
+                        mySum[i*strideF + j*hiddenSize + h] = out;
+                    }
+                    else if (packOutput == false and u + j < maxGLen)
+                        mySum[i*strideF + j*hiddenSize + h] = -1;
+                }
+            }
+            else if (packOutput == false and t + i < maxFLen){
+                // Again need to write finite data to don't-care region
+                #pragma unroll
+                for (int j = 0; j < tileG; ++j){
+                    if (u + j < maxGLen)
+                        mySum[i*strideF + j*hiddenSize + h] = -1;
+                }
+            }
+        }
+    }
+    else if (packOutput == false and t < maxFLen and u < maxGLen and hOffset+h < hiddenSize){
+        // Only need to ensure the finity in normal mode
+        #pragma unroll
+        for (int i = 0; i < tileF; ++i){
+            if (t + i < maxFLen){
+                #pragma unroll
+                for (int j = 0; j < tileG; ++j){
+                    if (u + j < maxGLen)
+                        mySum[i*strideF + j*hiddenSize + h] = -1;
+                }
+            }
+        }
+    }
+}
+
+/*
+Bwd operation (reduction) on one input tensor. Since the operation performed for the two input
+tensors are exactly the same, only one kernel is needed, and the different indexing offsets
+and strides are handled by OffsetCalBwd.
+
+When packing is enabled in the fwd op, unpacking is needed to restore the gradients in a 
+non-packed form.
+
+When ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,
+and mask contains the mask information.
+*/
+template <typename scalar_t, typename acc_t, class OffsetCal, bool masked>
+__device__ void transducer_joint_single_backward(
+    const scalar_t *grad,
+    const uint8_t *mask,
+    const int *fLen,
+    const int *gLen,
+    const int64_t *batchOffset,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    int64_t hiddenSize,
+    bool packOutput,
+    bool bwdFasterDim,  // whether bwd on the faster moving dimension (u)
+    float scale,
+    scalar_t *inGrad,
+    int yBlockOffset=0) {
+
+
+    const int batch = blockIdx.z;
+    // For the second input tensor, this offset need to be subtracted because the first yBlockOffset
+    // sets of thread blocks are for the first input tensor.
+    const int x = blockIdx.y-yBlockOffset;
+    const int hOffset = blockIdx.x*C10_WARP_SIZE;
+    const int wid = threadIdx.y;
+    const int lid = threadIdx.x;
+    const int numWarp = blockDim.y;
+    extern __shared__ char smem8[];
+    auto smem = reinterpret_cast<acc_t*>(smem8);
+
+    OffsetCal offsetCal(batch, batchOffset, fLen, gLen, maxFLen, maxGLen, hiddenSize, packOutput, 
+                        bwdFasterDim);
+    const auto maxXLen = offsetCal.getMaxXLen();
+    const auto myXLen = offsetCal.getMyXLen();
+    const auto myYLen = offsetCal.getMyYLen();
+    scalar_t *myInGrad = inGrad + batch*maxXLen*hiddenSize + x*hiddenSize + hOffset;
+    
+    if (x < myXLen){
+        
+        const auto myBatchOffset = offsetCal.getBatchOffset();
+        const auto strideX = offsetCal.getStrideX();
+        const auto strideY = offsetCal.getStrideY();
+        const scalar_t *myGrad = grad + myBatchOffset + x*strideX + hOffset;
+        const uint8_t *myMask = masked ? mask + myBatchOffset + x*strideX + hOffset : nullptr;
+        
+        // Each warp reduces numYPerWarp "y" first
+        acc_t warpSum = 0;
+        auto numYPerWarp = (myYLen+numWarp-1)/numWarp;
+        #pragma unroll
+        for (int warpY = 0; warpY < numYPerWarp; ++warpY){
+            auto y = wid*numYPerWarp + warpY;
+            if (y < myYLen and (hOffset+lid) < hiddenSize)
+                if (masked)
+                    warpSum += static_cast<acc_t>(myGrad[y*strideY + lid]) * myMask[y*strideY + lid] * scale;
+                else    
+                    warpSum += myGrad[y*strideY + lid];
+        }
+
+        // transpose partial sum in SMEM and reduce further using warpReduce
+        smem[lid*numWarp + wid] = warpSum;
+        __syncthreads();
+        auto sum = smem[wid*C10_WARP_SIZE + lid];
+        sum = warpReduce(sum, numWarp);
+
+        // a a b b c c d d
+        // a a b b c c d d
+        // a a b b c c d d
+        // a a b b c c d d
+        // example of 4 warps (a, b, c, d) with 8 threads per warp
+        // Each warp need 8 / 4 = 2 threads to write the results.
+        if (hOffset+wid*C10_WARP_SIZE/numWarp+lid/numWarp < hiddenSize){
+            if (lid % numWarp == 0){
+                myInGrad[wid*C10_WARP_SIZE/numWarp + lid/numWarp] = sum;
+            }
+        }
+    }
+    else if (wid == 0 and hOffset + lid < hiddenSize){
+        // Need to ensure the grad is zero for don't care region
+        myInGrad[lid] = 0;
+    }
+}
+
+/*
+Actual bwd (reduction) kernel get launched.
+Call transducer_joint_single_backward twice on two input tensors. 
+The two bwd ops are launched together, the first op uses blockIdx.y < maxFLen, and the second op 
+uses the rest.
+When ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,
+and mask contains the mask information.
+*/
+template <typename scalar_t, typename acc_t, class OffsetCal, bool masked>
+__global__ void transducer_joint_combined_backward(
+    const scalar_t *grad,
+    const uint8_t *mask,
+    const int *fLen,
+    const int *gLen,
+    const int64_t *batchOffset,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    int64_t hiddenSize,
+    bool packOutput,
+    float scale,
+    scalar_t *fGrad,
+    scalar_t *gGrad) {
+    if (blockIdx.y < maxFLen){
+        transducer_joint_single_backward<scalar_t, acc_t, OffsetCal, masked>(
+            grad,
+            mask,
+            fLen,
+            gLen,
+            batchOffset,
+            maxFLen,
+            maxGLen,
+            hiddenSize,
+            packOutput,
+            false,
+            scale,
+            fGrad);
+    }
+    else{
+        transducer_joint_single_backward<scalar_t, acc_t, OffsetCal, masked>(
+            grad,
+            mask,
+            fLen,
+            gLen,
+            batchOffset,
+            maxFLen,
+            maxGLen,
+            hiddenSize,
+            packOutput,
+            true,
+            scale,
+            gGrad,
+            maxFLen);
+    }  
+}
+
+/*
+Vectorized version of transducer_joint_single_backward
+Doing exact same operation as transducer_joint_single_backward except the load and store are
+vectorized.
+When packing is enabled in the fwd op, unpacking is needed to restore the gradients in a 
+non-packed form.
+When ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,
+and mask contains the mask information.
+*/
+template <typename scalar_t, typename acc_t, typename vec_t, int V, class OffsetCal, bool masked>
+__device__ void transducer_joint_single_vec_backward(
+    const scalar_t *grad,
+    const uint8_t *mask,
+    const int *fLen,
+    const int *gLen,
+    const int64_t *batchOffset,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    int64_t hiddenSize,
+    bool packOutput,
+    bool bwdFasterDim,
+    float scale,
+    scalar_t *inGrad,
+    int yBlockOffset=0){
+
+    const int batch = blockIdx.z;
+    const int x = blockIdx.y - yBlockOffset;
+    const int hOffset = blockIdx.x*C10_WARP_SIZE*V;
+    const int wid = threadIdx.y;
+    const int lid = threadIdx.x;
+    const int numWarp = blockDim.y;
+
+    // Figure out the vectorization type for mask
+    using mvec_t = mvec_type<V>;
+
+    OffsetCal offsetCal(batch, batchOffset, fLen, gLen, maxFLen, maxGLen, hiddenSize, packOutput, 
+                        bwdFasterDim);
+    const auto maxXLen = offsetCal.getMaxXLen();
+    const auto myXLen = offsetCal.getMyXLen();
+    const auto myYLen = offsetCal.getMyYLen();
+    scalar_t *myInGrad = inGrad + batch*maxXLen*hiddenSize + x*hiddenSize + hOffset;
+    extern __shared__ char smem8[];
+    auto smem = reinterpret_cast<acc_t*>(smem8);
+
+    acc_t warpSum[V];
+    scalar_t inBuffer[V];
+    uint8_t maskBuffer[V];
+    scalar_t outBuffer[V];
+    auto myInGradVec = reinterpret_cast<vec_t*>(myInGrad);
+    auto outBufferVec = reinterpret_cast<vec_t*>(outBuffer);
+
+    if (x < myXLen){
+        const auto myBatchOffset = offsetCal.getBatchOffset();
+        const auto strideX = offsetCal.getStrideX();
+        const auto strideY = offsetCal.getStrideY();
+        const scalar_t *myGrad = grad + myBatchOffset + x*strideX + hOffset;
+        const uint8_t *myMask = masked ? mask + myBatchOffset + x*strideX + hOffset
+                                            :nullptr;
+
+        for (int i = 0; i < V; ++i)
+            warpSum[i] = 0;
+
+        // Each warp reduces numYPerWarp "y" first
+        auto numYPerWarp = (myYLen+numWarp-1)/numWarp;
+        for (int warpY = 0; warpY < numYPerWarp; ++warpY){
+            auto y = wid*numYPerWarp + warpY;
+            auto myGradVec = reinterpret_cast<vec_t const *>(myGrad + y*strideY);
+            auto myMaskVec = masked ? reinterpret_cast<mvec_t const *>(myMask + y*strideY)
+                                        : nullptr;
+            auto inBufferVec = reinterpret_cast<vec_t*>(inBuffer);
+            auto maskBufferVec = reinterpret_cast<mvec_t*>(maskBuffer);
+            if (hOffset + lid*V < hiddenSize and y < myYLen){
+                *inBufferVec = myGradVec[lid];  // vectorized load
+                if (masked){
+                    *maskBufferVec = myMaskVec[lid];
+                    #pragma unroll
+                    for (int i = 0; i < V; ++i)
+                        warpSum[i] += static_cast<acc_t>(inBuffer[i]) * maskBuffer[i] * scale;
+                }
+                else{
+                    #pragma unroll
+                    for (int i = 0; i < V; ++i)
+                        warpSum[i] += inBuffer[i];
+                }
+            }
+        }
+        
+        // transpose partial sum in SMEM and reduce further using warpReduce
+        for (int i = 0; i < V; ++i){
+            smem[lid*numWarp + wid] = warpSum[i];
+            __syncthreads();
+            auto sum = smem[wid*C10_WARP_SIZE + lid];
+
+            if (hOffset+(wid*C10_WARP_SIZE/numWarp)*V < hiddenSize){
+                sum = warpReduce(sum, numWarp);
+                if (lid % numWarp == 0){
+                    outBuffer[i] = sum;
+                }
+            }
+            __syncthreads();
+        }
+
+        // a a b b c c d d
+        // a a b b c c d d
+        // a a b b c c d d
+        // a a b b c c d d
+        // example of 4 warps (a, b, c, d) with 8 threads per warp
+        // Each warp need 8 / 4 = 2 threads to write the results.
+        if (lid % numWarp == 0 and hOffset+(wid*C10_WARP_SIZE/numWarp + lid/numWarp)*V < hiddenSize)
+            myInGradVec[wid*C10_WARP_SIZE/numWarp + lid/numWarp] = *outBufferVec;     
+    }
+    else if (wid == 0 and hOffset + lid*V < hiddenSize){
+        // Need to ensure the grad is zero for don't care region
+        myInGradVec[lid] = 0;
+    }
+}
+
+/*
+Vecotrized version of transducer_joint_combined_backward
+Call transducer_joint_single_vec_backward twice on two input tensors. 
+The two bwd ops are launched together, the first op uses blockIdx.y < maxFLen, and the second op 
+uses the rest.
+When ReLU and/or dropout are performed in the fwd pass, this operation becomes a masked operation,
+and mask contains the mask information.
+*/
+template <typename scalar_t, typename acc_t, typename vec_t, int V, class OffsetCal, bool masked>
+__global__ void transducer_joint_combined_vec_backward(
+    const scalar_t *grad,
+    const uint8_t *mask,
+    const int *fLen,
+    const int *gLen,
+    const int64_t *batchOffset,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    int64_t hiddenSize,
+    bool packOutput,
+    float scale,
+    scalar_t *fGrad,
+    scalar_t *gGrad) {
+    if (blockIdx.y < maxFLen){
+        transducer_joint_single_vec_backward<scalar_t, acc_t, vec_t, V, OffsetCal, masked>(
+            grad,
+            mask,
+            fLen,
+            gLen,
+            batchOffset,
+            maxFLen,
+            maxGLen,
+            hiddenSize,
+            packOutput,
+            false,
+            scale,
+            fGrad);
+    }
+    else{
+        transducer_joint_single_vec_backward<scalar_t, acc_t, vec_t, V, OffsetCal, masked>(
+            grad,
+            mask,
+            fLen,
+            gLen,
+            batchOffset,
+            maxFLen,
+            maxGLen,
+            hiddenSize,
+            packOutput,
+            true,
+            scale,
+            gGrad,
+            maxFLen);
+    }  
+}
+
+
+
+
+std::vector<torch::Tensor> transducer_joint_cuda_forward(
+    torch::Tensor f,
+    torch::Tensor g,
+    torch::Tensor fLen,
+    torch::Tensor gLen,
+    torch::Tensor batchOffset,
+    int64_t packedBatch,
+    int opt,
+    bool packOutput,
+    bool relu,
+    bool dropout,
+    float dropoutProb,
+    int tileSize){
+
+    
+    auto tensorOpt = f.options();
+    auto dtype = f.scalar_type();
+    const auto batchSize = f.size(0);
+    const auto maxFLen = f.size(1);
+    const auto maxGLen = g.size(1);
+    const auto hiddenSize = f.size(2);
+    bool masked = dropout or relu;
+    
+    int64_t *batchOffsetPtr = nullptr;
+    torch::Tensor sum, mask;
+    auto maskOpt = tensorOpt.dtype(torch::kUInt8);
+    if (!packOutput){
+        sum = torch::empty({batchSize, maxFLen, maxGLen, hiddenSize}, tensorOpt);
+        batchOffsetPtr = nullptr;
+        if (masked)
+            mask = torch::empty({batchSize, maxFLen, maxGLen, hiddenSize}, maskOpt);
+    }
+    else{
+        sum = torch::empty({packedBatch, hiddenSize}, tensorOpt);    
+        batchOffsetPtr = batchOffset.data_ptr<int64_t>();
+        if (masked)
+            mask = torch::empty({packedBatch, hiddenSize}, maskOpt);
+    }
+    uint8_t *maskPtr = masked ? mask.data_ptr<uint8_t>() : nullptr;
+
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    TORCH_CHECK(opt == 0 or opt == 1, "Got an invalid optimization level ", opt);
+    // Simple heuristics
+    const int numThread = std::min(128, (static_cast<int>(hiddenSize)+C10_WARP_SIZE-1)
+                                        / C10_WARP_SIZE * C10_WARP_SIZE);
+    
+    if (opt == 0){
+        // vanilla kernel
+        const int threads = numThread;
+        const dim3 blocks(maxGLen, maxFLen, batchSize);
+
+        AT_DISPATCH_FLOATING_TYPES_AND_HALF(dtype, "transducer_joint_forward", ([&] {
+            transducer_joint_forward<scalar_t, OffsetCalFwd>
+            <<<blocks, threads, 0, stream>>>(
+                f.data_ptr<scalar_t>(), 
+                g.data_ptr<scalar_t>(), 
+                fLen.data_ptr<int>(), 
+                gLen.data_ptr<int>(), 
+                batchOffsetPtr,
+                maxFLen,
+                maxGLen,
+                hiddenSize,
+                packOutput,
+                sum.data_ptr<scalar_t>());
+        }));  
+    }
+    if (opt == 1){
+        // tiled version. For simplicity, assume tileF == tileG, even though the kernel can 
+        // support more general cases.
+        const int threads = numThread;
+        const int hiddenPerBlock = numThread;
+        const int hiddenBlock = (hiddenSize + hiddenPerBlock - 1) / hiddenPerBlock;
+        const dim3 blocks(  (maxGLen+tileSize-1)/tileSize * hiddenBlock, 
+                            (maxFLen+tileSize-1)/tileSize, 
+                            batchSize);
+
+        TORCH_CHECK(tileSize == 1 or tileSize == 2 or tileSize == 4, 
+                "Expected tileSize to be in [1, 2, 4], but got ", tileSize);
+
+        at::PhiloxCudaState rng_engine_inputs;
+        if (masked){
+            // set up PRG when the input is masked. rng_engine_inputs will be used as a space filler 
+            // for non-masked calls.
+            // Therefore no need to initialize.
+            c10::optional<at::Generator> gen_;
+            auto gen = at::get_generator_or_default<at::CUDAGeneratorImpl>(gen_, 
+                                                    at::cuda::detail::getDefaultCUDAGenerator());
+            // counterOffset records how many cuRAND calls each thread makes. For a tiled kernel, 
+            // each thread processes tileF * tileG output elements. 
+            int64_t counterOffset = tileSize * tileSize;
+            {
+                std::lock_guard<std::mutex> lock(gen->mutex_);
+                rng_engine_inputs = gen->philox_cuda_state(counterOffset);
+            }
+        }
+
+        AT_DISPATCH_FLOATING_TYPES_AND_HALF(dtype, "transducer_joint_forward", ([&] {
+            void(*kernel)(const scalar_t*, const scalar_t*, const int*, const int*, const int64_t*, 
+                            int64_t, int64_t, int64_t, int64_t, bool, bool, bool, float, 
+                            at::PhiloxCudaState, scalar_t*, uint8_t*);
+            if (masked){
+                switch (tileSize){
+                    case 2:
+                        kernel = &transducer_joint_tiled_forward<scalar_t, 2, 2, 4, OffsetCalFwd, 
+                                                                    true>;
+                        break;
+                    case 4:
+                        kernel = &transducer_joint_tiled_forward<scalar_t, 4, 4, 4, OffsetCalFwd, 
+                                                                    true>;
+                        break;
+                }
+            }
+            else{
+                switch (tileSize){
+                    case 1:
+                        kernel = &transducer_joint_tiled_forward<scalar_t, 1, 1, 4, OffsetCalFwd, 
+                                                                    false>;
+                        break;
+                    case 2:
+                        kernel = &transducer_joint_tiled_forward<scalar_t, 2, 2, 4, OffsetCalFwd, 
+                                                                    false>;
+                        break;
+                    case 4:
+                        kernel = &transducer_joint_tiled_forward<scalar_t, 4, 4, 4, OffsetCalFwd, 
+                                                                    false>;
+                        break;
+                }
+            }
+            
+            kernel<<<blocks, threads, 0, stream>>>(
+                f.data_ptr<scalar_t>(),
+                g.data_ptr<scalar_t>(),
+                fLen.data_ptr<int>(),
+                gLen.data_ptr<int>(),
+                batchOffsetPtr,
+                maxFLen,
+                maxGLen,
+                hiddenSize,
+                hiddenPerBlock,
+                packOutput,
+                relu,
+                dropout,
+                1.0f - dropoutProb,
+                rng_engine_inputs,
+                sum.data_ptr<scalar_t>(),
+                maskPtr);
+        }));  
+    }
+ 
+    C10_CUDA_CHECK(cudaGetLastError());
+    if (masked) 
+        return {sum, mask};
+    else
+        return {sum};
+}
+
+std::vector<torch::Tensor> transducer_joint_cuda_backward(
+    std::vector<torch::Tensor> in,
+    torch::Tensor fLen,
+    torch::Tensor gLen,
+    torch::Tensor batchOffset,
+    int maxFLen,
+    int maxGLen,
+    bool packOutput,
+    float scale){
+
+    auto grad = in[0];
+    bool masked = (in.size() == 2);
+    uint8_t *maskPtr = masked ? in[1].data_ptr<uint8_t>() : nullptr;
+
+    auto tensorOpt = grad.options();
+    auto dtype = grad.scalar_type();
+    const int batchSize = fLen.size(0);
+    const int hiddenSize = grad.size(-1);
+
+    const auto deviceProperties = at::cuda::getCurrentDeviceProperties();
+    const int maxNumWarp = deviceProperties->maxThreadsPerBlock / C10_WARP_SIZE;
+
+    torch::Tensor fGrad = torch::empty({batchSize, maxFLen, hiddenSize}, tensorOpt);
+    torch::Tensor gGrad = torch::empty({batchSize, maxGLen, hiddenSize}, tensorOpt);
+
+    int64_t *batchOffsetPtr = (!packOutput) ? nullptr : batchOffset.data_ptr<int64_t>(); 
+
+    // The number "y" I would like each thread to work on
+    const int workPerThread = 32;   
+    // Since the bwd for f and g have the same thread block size, we need to use the max of the two.
+    int numWarp = largestPowerOfTwo((std::max(maxFLen, maxGLen) + workPerThread-1) / workPerThread);
+    // Would like to have at least 2 warps 
+    numWarp = std::max(2, numWarp);
+    // cap on the maximum number of warps allowed
+    numWarp = std::min(maxNumWarp, numWarp); 
+
+    // Need smem for transposing the partial sum. The partial sum is in a matrix of the shape
+    // numWarp x warpSize
+    const int smemSize = numWarp * C10_WARP_SIZE;
+    const dim3 threads(C10_WARP_SIZE, numWarp, 1);
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(dtype, "transducer_joint_cuda_backward_kernel", ([&] {
+        auto gradPtr = grad.data_ptr<scalar_t>();
+        auto fLenPtr = fLen.data_ptr<int>();
+        auto gLenPtr = gLen.data_ptr<int>(); 
+        auto fGradPtr = fGrad.data_ptr<scalar_t>();
+        auto gGradPtr = gGrad.data_ptr<scalar_t>();
+
+        // resolve the acc_t type
+        using acc_t = at::acc_type<scalar_t, true>;
+        using vec_t = uint64_t;
+
+        constexpr int vectFactor = sizeof(vec_t) / sizeof(scalar_t);
+        constexpr int vecAlignment = std::alignment_of<vec_t>::value;
+
+        // if all input and output tensors meet the alignment requirement
+        bool memAlign = (reinterpret_cast<uint64_t>(gradPtr) % vecAlignment == 0) 
+                        and (reinterpret_cast<uint64_t>(fGradPtr) % vecAlignment == 0) 
+                        and (reinterpret_cast<uint64_t>(gGradPtr) % vecAlignment == 0);
+
+        if (vectFactor > 1 and hiddenSize%vectFactor == 0 and memAlign){
+            // If vectorization helps and the alignment requirement is met, use the vectorized 
+            // kernel. For simplicity, hiddenSize needs to be a multiple vecFactor.
+            const dim3 blocks(  (hiddenSize+C10_WARP_SIZE*vectFactor-1)/(C10_WARP_SIZE*vectFactor), 
+                                maxFLen+maxGLen, 
+                                batchSize);
+            if (masked){
+                transducer_joint_combined_vec_backward
+                    <scalar_t, acc_t, vec_t, vectFactor, OffsetCalBwd, true>
+                    <<<blocks, threads, smemSize*sizeof(acc_t)>>>(
+                    gradPtr,
+                    maskPtr,
+                    fLenPtr, 
+                    gLenPtr, 
+                    batchOffsetPtr, 
+                    maxFLen,
+                    maxGLen,
+                    hiddenSize,
+                    packOutput,
+                    scale,
+                    fGradPtr,
+                    gGradPtr);
+            }
+            else{
+                transducer_joint_combined_vec_backward
+                <scalar_t, acc_t, vec_t, vectFactor, OffsetCalBwd, false>
+                <<<blocks, threads, smemSize*sizeof(acc_t)>>>(
+                    gradPtr,
+                    maskPtr,
+                    fLenPtr, 
+                    gLenPtr, 
+                    batchOffsetPtr, 
+                    maxFLen,
+                    maxGLen,
+                    hiddenSize,
+                    packOutput,
+                    scale,
+                    fGradPtr,
+                    gGradPtr);    
+            }
+        }
+        else{
+            const dim3 blocks((hiddenSize+C10_WARP_SIZE-1)/C10_WARP_SIZE, 
+                                maxFLen + maxGLen, batchSize);
+            if (masked){
+                transducer_joint_combined_backward<scalar_t, acc_t, OffsetCalBwd, true>
+                <<<blocks, threads, smemSize*sizeof(acc_t)>>>(
+                    gradPtr,
+                    maskPtr,
+                    fLenPtr, 
+                    gLenPtr, 
+                    batchOffsetPtr, 
+                    maxFLen,
+                    maxGLen,
+                    hiddenSize,
+                    packOutput,
+                    scale,
+                    fGradPtr,
+                    gGradPtr);
+            }
+            else{
+                transducer_joint_combined_backward<scalar_t, acc_t, OffsetCalBwd, false>
+                <<<blocks, threads, smemSize*sizeof(acc_t)>>>(
+                    gradPtr,
+                    maskPtr,
+                    fLenPtr, 
+                    gLenPtr, 
+                    batchOffsetPtr, 
+                    maxFLen,
+                    maxGLen,
+                    hiddenSize,
+                    packOutput,
+                    scale,
+                    fGradPtr,
+                    gGradPtr);
+            }
+        }
+    }));   
+
+    return {fGrad, gGrad};
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_loss.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_loss.cpp
new file mode 100644
index 0000000000..f63a67f1e0
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_loss.cpp
@@ -0,0 +1,109 @@
+#include <torch/extension.h>
+#include <vector>
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<torch::Tensor> transducer_loss_cuda_forward(
+    torch::Tensor x,
+    torch::Tensor label,
+    torch::Tensor audLen,
+    torch::Tensor txtLen,
+    torch::Tensor batchOffset,
+    int maxFLen,
+    int blankIdx,
+    int opt,
+    bool packedInput);
+
+torch::Tensor transducer_loss_cuda_backward(
+    torch::Tensor x,
+    torch::Tensor lossGrad,
+    torch::Tensor alpha,
+    torch::Tensor beta,
+    torch::Tensor audLen,
+    torch::Tensor txtLen,
+    torch::Tensor label,
+    torch::Tensor batchOffset,
+    int maxFLen,
+    int blankIdx,
+    int opt,
+    bool fuseSoftmaxBackward,
+    bool packedInput);
+
+
+std::vector<torch::Tensor> transducer_loss_forward(
+    torch::Tensor x,
+    torch::Tensor label,
+    torch::Tensor fLen,
+    torch::Tensor yLen,
+    torch::Tensor batchOffset,
+    int maxFLen,
+    int blankIdx,
+    int opt,
+    bool packedInput
+    ) {
+
+    CHECK_INPUT(x);
+    CHECK_INPUT(label);
+    CHECK_INPUT(fLen);
+    CHECK_INPUT(yLen);
+    if (packedInput)
+        CHECK_INPUT(batchOffset);
+    return transducer_loss_cuda_forward(
+        x, 
+        label, 
+        fLen, 
+        yLen, 
+        batchOffset,
+        maxFLen,
+        blankIdx, 
+        opt,
+        packedInput);
+}
+
+torch::Tensor transducer_loss_backward(
+    torch::Tensor x,
+    torch::Tensor lossGrad,
+    torch::Tensor alpha,
+    torch::Tensor beta,
+    torch::Tensor fLen,
+    torch::Tensor yLen,
+    torch::Tensor label,
+    torch::Tensor batchOffset,
+    int maxFLen,
+    int blankIdx,
+    int opt,
+    bool fuseSoftmaxBackward,
+    bool packedInput){
+
+    CHECK_INPUT(x);
+    CHECK_INPUT(label);
+    CHECK_INPUT(lossGrad);
+    CHECK_INPUT(alpha);
+    CHECK_INPUT(beta);
+    CHECK_INPUT(fLen);
+    CHECK_INPUT(yLen);
+    if (packedInput)
+        CHECK_INPUT(batchOffset);
+
+    return transducer_loss_cuda_backward(
+        x,
+        lossGrad,
+        alpha,
+        beta,
+        fLen,
+        yLen,
+        label,
+        batchOffset,
+        maxFLen,
+        blankIdx,
+        opt,
+        fuseSoftmaxBackward,
+        packedInput);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &transducer_loss_forward, "transducer loss forward (CUDA)");
+  m.def("backward", &transducer_loss_backward, "transducer loss backward (CUDA)");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_loss_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_loss_kernel.cu
new file mode 100644
index 0000000000..295e14b3f1
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/transducer/transducer_loss_kernel.cu
@@ -0,0 +1,767 @@
+#include <vector>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <torch/extension.h>
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+
+template<typename scalar_t>
+__device__ __forceinline__ scalar_t logSumExp(scalar_t a, scalar_t b) {
+    // standard log-sum-exp trick is used here to provide better numerical stability
+    return (a >= b) ? a + std::log1p(exp(b-a)) : b + std::log1p(exp(a-b));
+}
+
+// Vanilla transducer loss function (i.e. forward-backward algorithm)
+// Detail of this loss function can be found in: 
+// [1] Sequence Transduction with Recurrent Neural Networks.
+
+// Forward (alpha) and backward (beta) path are launched together. Input is assumed to be converted
+// into log scale by the preceding log_softmax layer
+// Diagonal wavefront advancing usually used in dynamic programming is leveraged here. 
+// alpha and beta are of acc_t type, as they are essentially accumulators.
+
+// This loss function supports packed input where a tensor of shape [B, T, U, H] is packed into 
+// [B_packed, H].
+// Don't-care region (t > audLen) or (u > txtLen) is removed.
+// To support the packed input, the starting offsets for each batch need to be specified with
+// batchOffset.
+template <typename scalar_t, typename acc_t>
+__global__ void transducer_loss_forward(
+    const scalar_t* x,
+    const int* label,
+    const int* audLen,
+    const int* txtLen,
+    const int64_t* batchOffset,
+    int64_t dictSize,   // 64-bit indexing for data tensor
+    int64_t blankIdx,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    bool packedInput,
+    acc_t* alpha,
+    acc_t* beta,
+    scalar_t* loss) {
+
+    const int batch = blockIdx.y;
+    const int tid = threadIdx.x;
+    const auto myFLen = audLen[batch];
+    // Note that start of the sentence is added as 1 here
+    const auto myGLen = txtLen[batch] + 1;  
+    const auto myLabel = label + batch * (maxGLen-1);
+    const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch-1]) 
+                                                : batch * maxFLen * maxGLen;
+    const int64_t myStrideT = packedInput ? myGLen : maxGLen;
+    const scalar_t* myX = x + myBatchOffset * dictSize; 
+    int u  = tid;
+
+    if (blockIdx.x == 0){
+        // alpha path
+        acc_t* myAlpha = alpha + batch*maxFLen*maxGLen;
+        if (u == 0) 
+            myAlpha[0] = 0;
+        __syncthreads();
+
+        for (int64_t step = 1; step < myFLen+myGLen-1; ++step){
+            // Move along the diagonal wavefront to leverage available parallelism
+            for (u = tid; u < myGLen; u += blockDim.x){
+                int64_t t = step - u;
+                if (t >= 0 and t < myFLen and u >= 0 and u < myGLen){
+                    // Eq(16) in [1]
+                    if (u == 0){
+                        // alpha(t, u) = alpha(t-1, u) * null(t-1, u)
+                        myAlpha[t*maxGLen + u] = myAlpha[(t-1)*maxGLen] 
+                                                    + myX[((t-1)*myStrideT) * dictSize + blankIdx];
+                    }
+                    else if (t == 0){
+                        // alpha(t, u-1) = alpha(t, u-1) * y(t, u-1)
+                        myAlpha[u] = myAlpha[u - 1] + myX[(u - 1) * dictSize + myLabel[u - 1]];
+                    }
+                    else{
+                        // alpha(t, u) = alpha(t-1, u) * null(t-1, u) + alpha(t, u-1) * y(t, u-1)
+                        acc_t current = myAlpha[(t-1)*maxGLen + u] 
+                                        + myX[((t-1)*myStrideT + u) * dictSize + blankIdx];
+                        acc_t next = myAlpha[t*maxGLen + u - 1] 
+                                        + myX[(t*myStrideT + u - 1) * dictSize + myLabel[u - 1]];
+                        myAlpha[t*maxGLen + u] = logSumExp(next, current);
+                    }
+                }
+            }
+            __syncthreads();
+        }
+    }
+    else if (blockIdx.x == 1){
+        // beta path
+        acc_t* myBeta = beta + batch*maxFLen*maxGLen;
+        if (u == 0){
+            myBeta[(myFLen-1)*maxGLen + myGLen - 1] = myX[((myFLen-1)*myStrideT 
+                                                        + myGLen - 1) * dictSize + blankIdx];
+        }
+        __syncthreads();
+
+        for (int64_t step = myFLen+myGLen - 3; step >= 0; --step){
+            for (u = tid; u < myGLen; u += blockDim.x){
+                int64_t t = step - u;
+                if (t >= 0 and t < myFLen and u >=0 and u < myGLen){
+                    // Eq(18) in [1]
+                    if (u == myGLen - 1){
+                        // beta(t, u) = beta(t+1, u) * null(t, u)
+                        myBeta[t*maxGLen + u] = myBeta[(t+1)*maxGLen + u] 
+                                                + myX[(t*myStrideT + u) * dictSize + blankIdx];
+                    }
+                    else if (t == myFLen - 1){
+                        // beta(t, u) = beta(t, u+1) * y(t, u)
+                        myBeta[t*maxGLen + u] = myBeta[t*maxGLen + u + 1] 
+                                                + myX[(t*myStrideT + u) * dictSize + myLabel[u]];
+                    }
+                    else{
+                        // beta(t, u) = beta(t+1, u)*null(t, u) + beta(t, u+1)*y(t, u)
+                        acc_t current = myBeta[(t+1)*maxGLen + u] 
+                                        + myX[(t*myStrideT + u) * dictSize + blankIdx];
+                        acc_t next = myBeta[t*maxGLen + u + 1] 
+                                        + myX[(t*myStrideT + u) * dictSize + myLabel[u]];
+                        myBeta[t*maxGLen + u] = logSumExp(next, current);
+                    }
+                }
+            }
+            __syncthreads();
+        }
+        if (tid == 0)
+            loss[batch] = -myBeta[0];   
+    }
+
+}
+
+// transudcer loss function (i.e. forward-backward algorithm) with batch loading optimization.
+// Compared to the vanilla version, there are two optimizations:
+// 1. load x in batch through loop unrolling to reduce the latency.
+// 2. Use registers and shared memory to hold alpha and beta values passed from one step the next.
+// For simplicity, this kernel currently only supports U <= maxThread, which should be the common
+// case. For cases where U > maxThread, the vanilla kernel is used as a fallback option.
+
+// Detail of this loss function can be found in: 
+// [1] Sequence Transduction with Recurrent Neural Networks.
+// Forward (alpha) and backward (beta) path are launched together. Input is assumed to be converted
+// into log scale by the preceding log_softmax layer
+// Diagonal wavefront advancing usually used in dynamic programming is leveraged here.
+// alpha and beta are of acc_t type, as they are essentially accumulators.
+
+// This loss function supports packed input where a tensor of shape [B, T, U, H] is packed into 
+// [B_packed, H].
+// Don't-care region (t > audLen) or (u > txtLen) is removed.
+// To support the packed input, the starting offsets for each batch need to be specified with
+// batchOffset.
+template <typename scalar_t, typename acc_t, int batchLdSize>
+__global__ void transducer_loss_batch_load_forward(
+    const scalar_t* x,
+    const int* label,
+    const int* audLen,
+    const int* txtLen,
+    const int64_t* batchOffset,
+    int64_t dictSize,
+    int64_t blankIdx,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    bool packedInput,
+    acc_t* alpha,
+    acc_t* beta,
+    scalar_t* loss) {
+
+    const int batch = blockIdx.y;
+    int u  = threadIdx.x;
+    const auto myFLen = audLen[batch];
+    const auto myGLen = txtLen[batch] + 1;
+    const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch-1]) 
+                                                : batch * maxFLen * maxGLen;
+    const int64_t myStrideT = packedInput ? myGLen : maxGLen;
+    const scalar_t* myX = x + myBatchOffset * dictSize; 
+    scalar_t next[batchLdSize], current[batchLdSize];
+    extern __shared__ char smem8[];
+    auto smem = reinterpret_cast<acc_t*>(smem8);
+
+    if (blockIdx.x == 0){
+        // alpha path
+        acc_t* myAlpha = alpha + batch*maxFLen*maxGLen;
+        // two SMEM regions for double buffering read and write data to avoid data race
+        acc_t * const sharedAlpha[2] = {smem, smem+maxGLen};
+
+        sharedAlpha[0][u] = 0; 
+        __syncthreads();
+
+        if (u == 0)
+            myAlpha[0] = 0;
+
+        auto myAlphaLabel = (u == 0) ? 0 : label[batch*(maxGLen-1) + u - 1];
+        // register used to pass value to the next step for the same thread
+        acc_t prvStepAlpha = 0;
+        for (int64_t step = 1; step < myFLen+myGLen-1+batchLdSize; step += batchLdSize){
+            // Move along the diagonal wavefront to leverage available parallelism
+            // Batch loading X through loop unrolling
+            #pragma unroll
+            for (int i = 0; i < batchLdSize; ++i){
+                if (step+i<myFLen+myGLen-1){
+                    // index computing
+                    int64_t t = step + i - u;
+                    int64_t currentId = ((t-1)*myStrideT + u) * dictSize + blankIdx;
+                    int64_t nextId = (t*myStrideT + u - 1) * dictSize + myAlphaLabel;
+                    // main loading loop
+                    if (t >= 0 and t < myFLen and u >= 0 and u < myGLen){
+                        if (u == 0){
+                            current[i] = myX[currentId];
+                        }
+                        else if (t == 0){
+                            next[i] = myX[nextId];
+                        }
+                        else{
+                            current[i] = myX[currentId];
+                            next[i] = myX[nextId];
+                        }
+                    }
+                }
+            }
+            // main computing loop
+            for (int i = 0; i < batchLdSize; ++i){
+                // swap the pointer for double buffering
+                auto sharedAlphaRd = sharedAlpha[(step+i-1)%2];
+                auto sharedAlphaWr = sharedAlpha[(step+i)%2];
+                if (step+i<myFLen+myGLen-1){
+                    int64_t t = step + i - u;
+                    if (t >= 0 and t < myFLen and u >= 0 and u < myGLen){
+                        // Eq(16) in [1]
+                        if (u == 0)
+                            prvStepAlpha = prvStepAlpha+current[i];
+                        else if (t == 0)
+                            prvStepAlpha = sharedAlphaRd[u-1]+next[i];
+                        else
+                            prvStepAlpha = logSumExp(prvStepAlpha+current[i], sharedAlphaRd[u-1]
+                                            + next[i]);
+                        sharedAlphaWr[u] = prvStepAlpha;
+                        myAlpha[t*maxGLen + u] = prvStepAlpha;
+                    }
+                }
+                __syncthreads();
+            }
+        }
+    }
+    else if (blockIdx.x == 1){
+        // beta path
+        acc_t* myBeta = beta + batch*maxFLen*maxGLen;
+        // two SMEM regions for double buffering read and write data to avoid data race
+        acc_t * const sharedBeta[2] = {smem, smem + maxGLen};
+        sharedBeta[0][u] = myX[((myFLen-1)*myStrideT + myGLen - 1) * dictSize + blankIdx];
+        __syncthreads();
+
+        auto myBetaLabel = (u == maxGLen - 1) ? 0 : label[batch*(maxGLen-1) + u];
+        // register used to pass value to the next step for the same thread
+        acc_t prvStepBeta = myX[((myFLen-1)*myStrideT + myGLen - 1) * dictSize + blankIdx];
+        if (u == 0)
+            myBeta[(myFLen-1)*maxGLen + myGLen - 1] = prvStepBeta;
+
+        for (int64_t step = 1; step < myFLen+myGLen-1; step += batchLdSize){
+            // Move along the diagonal wavefront to leverage available parallelism
+            // Batch loading X
+            #pragma unroll
+            for (int i = 0; i < batchLdSize; ++i){
+                if (step+i<myFLen+myGLen-1){
+                    // index computing
+                    int64_t t = myFLen+myGLen - (step + i) - 2 - u;
+                    int64_t currentId = (t*myStrideT + u) * dictSize + blankIdx;
+                    int64_t nextId = (t*myStrideT + u) * dictSize + myBetaLabel;
+                    // main loading loop
+                    if (t >= 0 and t < myFLen and u >= 0 and u < myGLen){
+                        if (u == myGLen - 1){
+                            current[i] = myX[currentId];
+                        }
+                        else if (t == myFLen - 1){
+                            next[i] = myX[nextId];
+                        }
+                        else{
+                            current[i] = myX[currentId];
+                            next[i] = myX[nextId];
+                        }
+                    }
+                }
+            }
+            // main computing loop
+            for (int i = 0; i < batchLdSize; ++i){
+                // swap the pointer for double buffering
+                auto sharedBetaRd = sharedBeta[(step+i-1)%2];
+                auto sharedBetaWr = sharedBeta[(step+i)%2];
+                if (step+i<myFLen+myGLen-1){
+                    int64_t t = myFLen+myGLen - (step + i) - 2 - u;
+                    if (t >= 0 and t < myFLen and u >= 0 and u < myGLen){
+                        // Eq(18) in [1]
+                        if (u == myGLen - 1)
+                            prvStepBeta = prvStepBeta+current[i];
+                        else if (t == myFLen - 1)
+                            prvStepBeta = sharedBetaRd[u+1]+next[i];
+                        else
+                            prvStepBeta = logSumExp(prvStepBeta+current[i], sharedBetaRd[u+1]
+                                            + next[i]);
+                        sharedBetaWr[u] = prvStepBeta;
+                        myBeta[t*maxGLen + u] = prvStepBeta;
+                    }
+                    
+                }
+                __syncthreads();
+            }
+        }
+        if (u == 0)
+            loss[batch] = -prvStepBeta; 
+    }
+
+}
+
+// Vanilla transudcer loss backward operation.
+// Detail of this loss function can be found in: 
+// [1] Sequence Transduction with Recurrent Neural Networks.
+// For this backward kernel, bwd op for the preceding softmax is assumed to be handled elsewhere, 
+// hence only Eq(20) in [1] is implemented in this kernel.
+
+// Each thread block works on [batch, t, :, :] of data. Each thread works on a specific u at a time
+// Since only gradients for the correct token and null token need to be updated, gradients at other
+// locations are initialized to 0.
+
+// To support the packed input, the starting offsets for each batch need to be specified with
+// batchOffset.
+template <typename scalar_t, typename acc_t>
+__global__ void transducer_loss_backward(
+    const scalar_t* x,
+    const scalar_t* lossGrad,
+    const int* audLen,
+    const int* txtLen,
+    const int* label,
+    const acc_t* alpha,
+    const acc_t* beta,
+    const int64_t* batchOffset,
+    int64_t dictSize,
+    int64_t blankIdx,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    bool packedInput,
+    scalar_t* xGrad) {
+
+    const int tid = threadIdx.x;
+    const int t = blockIdx.x;
+    const int batch = blockIdx.y;
+    const int64_t myFLen = audLen[batch];
+    const int64_t myGLen = txtLen[batch] + 1;
+    const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch-1]) 
+                                                : batch * maxFLen * maxGLen;
+    const int64_t myStrideT = packedInput ? myGLen : maxGLen;
+    auto myX = x + (myBatchOffset + t*myStrideT)*dictSize;
+    auto myAlpha = alpha + batch*maxFLen*maxGLen;
+    auto myBeta = beta + batch*maxFLen*maxGLen;
+    auto myXGrad = xGrad + (myBatchOffset + t*myStrideT)*dictSize; 
+    auto myLabel = label + batch*(maxGLen-1);
+
+    int64_t u = tid;
+    while (t < myFLen and u < myGLen){
+        // Do the update
+        // loss = -ln(Pr(y*|x))
+        acc_t grad = std::log(lossGrad[batch]) + myAlpha[t*maxGLen + u] - myBeta[0];  
+        if (u != myGLen - 1)
+            myXGrad[u*dictSize + myLabel[u]] = -std::exp(grad + myBeta[t*maxGLen + u + 1] 
+                                                + myX[u*dictSize + myLabel[u]]);
+        if (t == myFLen - 1 and u == myGLen - 1)
+            myXGrad[u*dictSize + blankIdx] = -std::exp(grad + myX[u*dictSize + blankIdx]);
+        else if (t != myFLen - 1)
+            myXGrad[u*dictSize + blankIdx] = -std::exp(grad + myBeta[(t+1)*maxGLen + u] 
+                                                + myX[u*dictSize + blankIdx]); 
+
+        u += blockDim.x;
+    }
+}
+
+// Fused transudcer loss backward operation.
+// Detail of this loss function can be found in: 
+// [1] Sequence Transduction with Recurrent Neural Networks.
+// The bwd op of the preceding softmax layer is fused in this kernel. 
+// Each thread block works on [batch, t, u, :] of data. Each thread works on a specific h at a time
+
+// To support the packed input, the starting offsets for each batch need to be specified with
+// batchOffset.
+template <typename scalar_t, typename acc_t>
+__global__ void transducer_loss_fused_backward(
+    const scalar_t* x,
+    const scalar_t* lossGrad,
+    const int* audLen,
+    const int* txtLen,
+    const int* label,
+    const acc_t* alpha,
+    const acc_t* beta,
+    const int64_t* batchOffset,
+    int64_t dictSize,
+    int64_t blankIdx,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    bool packedInput,
+    scalar_t* xGrad) {
+    
+    const int tid = threadIdx.x;
+    const int u = blockIdx.x;
+    const int t = blockIdx.y;
+    const int batch = blockIdx.z;
+    const int64_t myFLen = audLen[batch];
+    const int64_t myGLen = txtLen[batch] + 1;
+    const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch-1]) 
+                                                : batch * maxFLen * maxGLen;
+    const int64_t myStrideT = packedInput ? myGLen : maxGLen;
+
+    __shared__ acc_t commonFactor, myBetaTU, myBetaTUp1, myBetaTp1U, myLabelShared;
+    auto myXGrad = xGrad + (myBatchOffset + t*myStrideT +u)*dictSize; 
+
+    if (t < myFLen and u < myGLen){ 
+        auto myX = x + (myBatchOffset + t*myStrideT +u)*dictSize; 
+        auto myAlpha = alpha + batch*maxFLen*maxGLen;
+        auto myBeta = beta + batch*maxFLen*maxGLen;
+        auto myLabel = label + batch*(maxGLen-1);
+
+        // load and store shared variables in SMEM
+        if (tid == 0){
+            commonFactor = std::log(lossGrad[batch]) + myAlpha[t*maxGLen + u] - myBeta[0];
+            myBetaTU = myBeta[t*maxGLen + u];
+            myBetaTUp1 = myBeta[t*maxGLen + u + 1];
+            myBetaTp1U = myBeta[(t+1)*maxGLen + u];
+            myLabelShared = myLabel[u];
+        }
+
+        __syncthreads();
+
+        for (int64_t h = tid; h < dictSize; h += blockDim.x){
+            // Do the update
+            acc_t grad = commonFactor + myX[h];  // loss = -ln(Pr(y*|x))
+            acc_t myGrad = std::exp(grad + myBetaTU);
+            if (u != myGLen - 1 and h == myLabelShared){
+                myGrad -= std::exp(grad + myBetaTUp1);
+            }
+            else if (h == blankIdx){
+                if (t == myFLen - 1 and u == myGLen - 1)
+                    myGrad -= std::exp(grad);
+                else if (t != myFLen - 1)
+                    myGrad -= std::exp(grad + myBetaTp1U);
+            }
+            myXGrad[h] = myGrad;
+        }
+    }
+    else if (!packedInput){
+        // In non-pack mode, need to make sure the gradients for don't-care regions are zero.
+        for (int64_t h = tid; h < dictSize; h += blockDim.x){
+            myXGrad[h] = 0;
+        }
+    }
+}
+
+
+// Vectorized version of fused transudcer loss backward operation.
+// Detail of this loss function can be found in: 
+// [1] Sequence Transduction with Recurrent Neural Networks.
+// The bwd op of the preceding softmax layer is fused in this kernel. 
+// Each thread block works on [batch, t, u, :] of data. Each thread works on a specific h at a time
+
+// To support the packed input, the starting offsets for each batch need to be specified with
+// batchOffset.
+template <typename scalar_t, typename acc_t, typename vec_t, int V>
+__global__ void transducer_loss_fused_vec_backward(
+    const scalar_t* x,
+    const scalar_t* lossGrad,
+    const int* audLen,
+    const int* txtLen,
+    const int* label,
+    const acc_t* alpha,
+    const acc_t* beta,
+    const int64_t* batchOffset,
+    int64_t dictSize,
+    int64_t blankIdx,
+    int64_t maxFLen,
+    int64_t maxGLen,
+    bool packedInput,
+    scalar_t* xGrad) {
+    
+    const int tid = threadIdx.x;
+    const int u = blockIdx.x;
+    const int t = blockIdx.y;
+    const int batch = blockIdx.z;
+    const int64_t myFLen = audLen[batch];
+    const int64_t myGLen = txtLen[batch] + 1;
+    const int64_t myBatchOffset = packedInput ? (batch == 0 ? 0 : batchOffset[batch-1]) 
+                                                : batch * maxFLen * maxGLen;
+    const int64_t myStrideT = packedInput ? myGLen : maxGLen;
+
+    __shared__ acc_t commonFactor, myBetaTU, myBetaTUp1, myBetaTp1U, myLabelShared;
+    auto myXGrad = xGrad + (myBatchOffset + t*myStrideT +u)*dictSize; 
+    auto myX = x + (myBatchOffset + t*myStrideT +u)*dictSize; 
+    auto myAlpha = alpha + batch*maxFLen*maxGLen;
+    auto myBeta = beta + batch*maxFLen*maxGLen;
+    auto myLabel = label + batch*(maxGLen-1);
+
+    // Variabels for vectorization
+    scalar_t myXBuffer[V], myXGradBuffer[V];
+    auto myXVec = reinterpret_cast<vec_t const *>(myX);
+    auto myXGradVec = reinterpret_cast<vec_t*>(myXGrad);
+    auto myXBufferVec = reinterpret_cast<vec_t*>(myXBuffer);
+    auto myXGradBufferVec = reinterpret_cast<vec_t*>(myXGradBuffer);
+    if (t < myFLen and u < myGLen){ 
+        // load and store shared variables in SMEM
+        if (tid == 0){
+            commonFactor = std::log(lossGrad[batch]) + myAlpha[t*maxGLen + u] - myBeta[0];
+            myBetaTU = myBeta[t*maxGLen + u];
+            if (t != myFLen - 1)
+                myBetaTp1U = myBeta[(t+1)*maxGLen + u];
+            if (u != myGLen - 1){
+                myBetaTUp1 = myBeta[t*maxGLen + u + 1];
+                myLabelShared = myLabel[u];
+            }
+        }
+
+        __syncthreads();
+
+        #pragma unroll
+        for (int64_t h0 = tid*V; h0 < dictSize; h0 += blockDim.x*V){
+            // Load myX in a vector form
+            *myXBufferVec = myXVec[h0/V];
+            // Do the update for a vector of input
+            #pragma unroll
+            for (int i = 0; i < V; ++i){
+                auto h = h0 + i;
+                acc_t grad = commonFactor + myXBuffer[i];  // loss = -ln(Pr(y*|x))
+                acc_t myGrad = std::exp(grad + myBetaTU);
+                if (u != myGLen - 1 and h == myLabelShared){
+                    myGrad -= std::exp(grad + myBetaTUp1);
+                }
+                else if (h == blankIdx){
+                    if (t == myFLen - 1 and u == myGLen - 1)
+                        myGrad -= std::exp(grad);
+                    else if (t != myFLen - 1)
+                        myGrad -= std::exp(grad + myBetaTp1U);
+                }
+                myXGradBuffer[i] = myGrad;
+            }
+
+            // Store myXGrad in a vector form
+            myXGradVec[h0/V] = *myXGradBufferVec;
+            
+        }
+    }
+    else if (!packedInput){
+        // In non-pack mode, need to make sure the gradients for don't-care regions are zero.
+        for (int64_t h0 = tid*V; h0 < dictSize; h0 += blockDim.x*V){
+            myXGradVec[h0/V] = 0;
+        }
+    }
+}
+
+
+std::vector<torch::Tensor> transducer_loss_cuda_forward(
+    torch::Tensor x,
+    torch::Tensor label,
+    torch::Tensor audLen,
+    torch::Tensor txtLen,
+    torch::Tensor batchOffset,
+    int maxFLen,
+    int blankIdx,
+    int opt,
+    bool packedInput){
+
+    auto scalarType = x.scalar_type();
+    auto tensorOpt = x.options();
+    const int batchSize = label.size(0);
+    const int maxGLen = label.size(1) + 1;
+    const int dictSize = x.size(-1);
+
+    TORCH_CHECK(blankIdx >= 0 and blankIdx < dictSize, 
+                "Expected blank index to be in the range of 0 to ", 
+                dictSize-1,
+                ", but got ", 
+                blankIdx);
+    TORCH_CHECK(opt == -1 or opt == 0 or opt == 1, 
+                "Got an invalid optimization level ", 
+                opt);
+
+    // The data type of alpha and beta will be resolved at dispatch time,
+    // hence defined here and assigned later
+    torch::Tensor alpha;    
+    torch::Tensor beta;
+    torch::Tensor loss = torch::empty({batchSize}, tensorOpt);
+    const auto deviceProperties = at::cuda::getCurrentDeviceProperties();
+    const auto maxThreadPerBlock = deviceProperties->maxThreadsPerBlock;
+    const auto maxSmemPerBlock = deviceProperties->sharedMemPerBlock;
+    const auto batchOffsetPtr = packedInput ? batchOffset.data_ptr<int64_t>() : nullptr;
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+
+    AT_DISPATCH_FLOATING_TYPES_AND_HALF(scalarType, "transducer_loss_cuda_forward", ([&] {
+        // resolve accumulation type
+        using acc_t = at::acc_type<scalar_t, true>;
+        auto accType = c10::CppTypeToScalarType<acc_t>::value;
+        auto accTensorOpt = tensorOpt.dtype(accType);
+        alpha = torch::empty({batchSize, maxFLen, maxGLen}, accTensorOpt);
+        beta = torch::empty({batchSize, maxFLen, maxGLen}, accTensorOpt);
+
+        // decide what kernel to launch based on the problem size
+        // if the required SMEM size or number threads exceeds the limit, fall back to the vanilla
+        // kernel.
+        const auto smemSize = 2*maxGLen*sizeof(acc_t);
+        const auto optFallBack = (maxGLen > maxThreadPerBlock or smemSize > maxSmemPerBlock) ? 0 
+                                    : (opt == -1) ? 1 : opt;
+        const int threads = std::min(maxThreadPerBlock, maxGLen);
+        const dim3 blocks(2, batchSize, 1);        
+
+        if (optFallBack == 0)
+            transducer_loss_forward<<<blocks, threads, 0, stream>>>(
+                x.data_ptr<scalar_t>(), 
+                label.data_ptr<int>(), 
+                audLen.data_ptr<int>(), 
+                txtLen.data_ptr<int>(), 
+                batchOffsetPtr,
+                dictSize, 
+                blankIdx, 
+                maxFLen,
+                maxGLen,
+                packedInput,
+                alpha.data_ptr<acc_t>(), 
+                beta.data_ptr<acc_t>(), 
+                loss.data_ptr<scalar_t>());
+        else if (optFallBack == 1)
+            transducer_loss_batch_load_forward<scalar_t, acc_t, 4>
+            <<<blocks, threads, smemSize, stream>>>(
+                x.data_ptr<scalar_t>(), 
+                label.data_ptr<int>(), 
+                audLen.data_ptr<int>(), 
+                txtLen.data_ptr<int>(), 
+                batchOffsetPtr,
+                dictSize, 
+                blankIdx, 
+                maxFLen,
+                maxGLen,
+                packedInput,
+                alpha.data_ptr<acc_t>(), 
+                beta.data_ptr<acc_t>(), 
+                loss.data_ptr<scalar_t>());  
+
+    }));
+    C10_CUDA_CHECK(cudaGetLastError());
+
+    return {alpha, beta, loss};
+}
+
+
+
+
+torch::Tensor transducer_loss_cuda_backward(
+    torch::Tensor x,
+    torch::Tensor lossGrad,
+    torch::Tensor alpha,
+    torch::Tensor beta,
+    torch::Tensor audLen,
+    torch::Tensor txtLen,
+    torch::Tensor label,
+    torch::Tensor batchOffset,
+    int maxFLen,
+    int blankIdx,
+    int opt,
+    bool fuseSoftmaxBackward,
+    bool packedInput){
+
+    auto dtype = x.scalar_type();
+    torch::Tensor xGrad;
+    const int batchSize = label.size(0);
+    const int maxGLen = label.size(1) + 1;
+    const int dictSize = x.size(-1);
+    const auto deviceProperties = at::cuda::getCurrentDeviceProperties();
+    const int maxThreadPerBlock = deviceProperties->maxThreadsPerBlock;
+    const int warpSize = deviceProperties->warpSize;
+    const auto batchOffsetPtr = packedInput ? batchOffset.data_ptr<int64_t>() : nullptr;
+    cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+    
+    if (fuseSoftmaxBackward){
+        // alloc empty tensors for performance, hence need to ensure zeros are writtern to 
+        // don't-care region in the kernel.
+        xGrad = torch::empty_like(x);
+
+        // Would like each thread to work on 4 hidden units
+        const int workPerThread = 4;  
+        // Don't want to have more than 128 threads per thread block
+        const int maxThreadPerElmt = std::min(128, maxThreadPerBlock);
+        const int threads = std::min(maxThreadPerElmt, std::max(warpSize, 
+                                    (dictSize+workPerThread-1)/workPerThread));
+        const dim3 blocks(maxGLen, maxFLen, batchSize);
+
+        AT_DISPATCH_FLOATING_TYPES_AND_HALF(dtype, "transducer_loss_cuda_backward", ([&] {
+            using vec_t = uint64_t;
+            using acc_t = at::acc_type<scalar_t, true>;
+            constexpr int vectFactor = sizeof(vec_t) / sizeof(scalar_t);
+            constexpr int vecAlignment = std::alignment_of<vec_t>::value;
+            // if all input and output tensors meet the alignment requirement
+            bool memAlign = reinterpret_cast<uint64_t>(x.data_ptr<scalar_t>()) % vecAlignment == 0
+                                and reinterpret_cast<uint64_t>(xGrad.data_ptr<scalar_t>()) 
+                                        % vecAlignment == 0;
+
+            if (vectFactor > 1 and dictSize%vectFactor == 0 and memAlign){
+                transducer_loss_fused_vec_backward<scalar_t, acc_t, vec_t, vectFactor>
+                    <<<blocks, threads, 0, stream>>>(    
+                    x.data_ptr<scalar_t>(), 
+                    lossGrad.data_ptr<scalar_t>(),
+                    audLen.data_ptr<int>(), 
+                    txtLen.data_ptr<int>(), 
+                    label.data_ptr<int>(),
+                    alpha.data_ptr<acc_t>(), 
+                    beta.data_ptr<acc_t>(),  
+                    batchOffsetPtr,
+                    dictSize, 
+                    blankIdx, 
+                    maxFLen,
+                    maxGLen,
+                    packedInput,
+                    xGrad.data_ptr<scalar_t>());   
+            }
+            else{
+                transducer_loss_fused_backward<<<blocks, threads, 0, stream>>>(    
+                    x.data_ptr<scalar_t>(), 
+                    lossGrad.data_ptr<scalar_t>(),
+                    audLen.data_ptr<int>(), 
+                    txtLen.data_ptr<int>(), 
+                    label.data_ptr<int>(),
+                    alpha.data_ptr<acc_t>(), 
+                    beta.data_ptr<acc_t>(),  
+                    batchOffsetPtr,
+                    dictSize, 
+                    blankIdx, 
+                    maxFLen,
+                    maxGLen,
+                    packedInput,
+                    xGrad.data_ptr<scalar_t>());   
+                
+            }
+        }));
+    }
+    else{
+        // for non-fused kernel, the gradients need to be writtern are very sparse, hence initialize
+        // the tensor with all zeros.
+        xGrad = torch::zeros_like(x);
+        // don't launch more threads than needed.
+        const int threads = std::min(maxThreadPerBlock, maxGLen);
+        const dim3 blocks(maxFLen, batchSize);
+        AT_DISPATCH_FLOATING_TYPES_AND_HALF(dtype, "transducer_loss_cuda_backward", ([&] {
+            using acc_t = at::acc_type<scalar_t, true>;
+            transducer_loss_backward<<<blocks, threads, 0, stream>>>(    
+                x.data_ptr<scalar_t>(), 
+                lossGrad.data_ptr<scalar_t>(),
+                audLen.data_ptr<int>(), 
+                txtLen.data_ptr<int>(), 
+                label.data_ptr<int>(),
+                alpha.data_ptr<acc_t>(), 
+                beta.data_ptr<acc_t>(), 
+                batchOffsetPtr, 
+                dictSize, 
+                blankIdx, 
+                maxFLen,
+                maxGLen,
+                packedInput,
+                xGrad.data_ptr<scalar_t>());
+        }));
+    }
+    C10_CUDA_CHECK(cudaGetLastError());
+    
+    return xGrad;
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/xentropy/interface.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/xentropy/interface.cpp
new file mode 100644
index 0000000000..3a8c201ca9
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/xentropy/interface.cpp
@@ -0,0 +1,52 @@
+#include <torch/extension.h>
+
+// CUDA forward declarations
+
+std::vector<at::Tensor> softmax_xentropy_cuda(
+    const at::Tensor &input,
+    const at::Tensor &labels,
+    const float smoothing,
+    const bool half_to_float);
+
+at::Tensor softmax_xentropy_backward_cuda(
+    const at::Tensor &grad_loss,
+    const at::Tensor &logits,
+    const at::Tensor &max_log_sum_exp,
+    const at::Tensor &labels,
+    const float smoothing);
+
+// C++ interface
+
+#define CHECK_CUDA(x) AT_ASSERTM(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) AT_ASSERTM(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<at::Tensor> softmax_xentropy_forward(
+    const at::Tensor &input,
+    const at::Tensor &labels,
+    const float smoothing,
+    const bool half_to_float) {
+    CHECK_CUDA(input);
+    CHECK_INPUT(labels);
+
+    return softmax_xentropy_cuda(input, labels, smoothing, half_to_float);
+}
+
+at::Tensor softmax_xentropy_backward(
+    const at::Tensor &grad_loss,
+    const at::Tensor &logits,
+    const at::Tensor &max_log_sum_exp,
+    const at::Tensor &labels,
+    const float smoothing)  {
+    CHECK_CUDA(grad_loss);
+    CHECK_CUDA(logits);
+    CHECK_INPUT(max_log_sum_exp);
+    CHECK_INPUT(labels);
+
+    return softmax_xentropy_backward_cuda(grad_loss, logits, max_log_sum_exp, labels, smoothing);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+    m.def("forward", &softmax_xentropy_forward, "Softmax cross entropy loss with label smoothing forward (CUDA)");
+    m.def("backward", &softmax_xentropy_backward, "Softmax cross entropy loss with label smoothing backward (CUDA)");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/xentropy/xentropy_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/xentropy/xentropy_kernel.cu
new file mode 100644
index 0000000000..7f4190dd15
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/csrc/xentropy/xentropy_kernel.cu
@@ -0,0 +1,718 @@
+/**
+ * From PyTorch:
+ *
+ * Copyright (c) 2016-     Facebook, Inc            (Adam Paszke)
+ * Copyright (c) 2014-     Facebook, Inc            (Soumith Chintala)
+ * Copyright (c) 2011-2014 Idiap Research Institute (Ronan Collobert)
+ * Copyright (c) 2012-2014 Deepmind Technologies    (Koray Kavukcuoglu)
+ * Copyright (c) 2011-2012 NEC Laboratories America (Koray Kavukcuoglu)
+ * Copyright (c) 2011-2013 NYU                      (Clement Farabet)
+ * Copyright (c) 2006-2010 NEC Laboratories America (Ronan Collobert, Leon Bottou, Iain Melvin, Jason Weston)
+ * Copyright (c) 2006      Idiap Research Institute (Samy Bengio)
+ * Copyright (c) 2001-2004 Idiap Research Institute (Ronan Collobert, Samy Bengio, Johnny Mariethoz)
+ *
+ * From Caffe2:
+ *
+ * Copyright (c) 2016-present, Facebook Inc. All rights reserved.
+ *
+ * All contributions by Facebook:
+ * Copyright (c) 2016 Facebook Inc.
+ *
+ * All contributions by Google:
+ * Copyright (c) 2015 Google Inc.
+ * All rights reserved.
+ *
+ * All contributions by Yangqing Jia:
+ * Copyright (c) 2015 Yangqing Jia
+ * All rights reserved.
+ *
+ * All contributions from Caffe:
+ * Copyright(c) 2013, 2014, 2015, the respective contributors
+ * All rights reserved.
+ *
+ * All other contributions:
+ * Copyright(c) 2015, 2016 the respective contributors
+ * All rights reserved.
+ *
+ * Caffe2 uses a copyright model similar to Caffe: each contributor holds
+ * copyright over their contributions to Caffe2. The project versioning records
+ * all such contribution and copyright details. If a contributor wants to further
+ * mark their specific copyright on a particular contribution, they should
+ * indicate their copyright solely in the commit message of the change when it is
+ * committed.
+ *
+ * All rights reserved.
+ *
+ * Redistribution and use in source and binary forms, with or without
+ * modification, are permitted provided that the following conditions are met:
+ *
+ * 1. Redistributions of source code must retain the above copyright
+ *    notice, this list of conditions and the following disclaimer.
+ *
+ * 2. Redistributions in binary form must reproduce the above copyright
+ *    notice, this list of conditions and the following disclaimer in the
+ *    documentation and/or other materials provided with the distribution.
+ *
+ * 3. Neither the names of Facebook, Deepmind Technologies, NYU, NEC Laboratories America
+ *    and IDIAP Research Institute nor the names of its contributors may be
+ *    used to endorse or promote products derived from this software without
+ *    specific prior written permission.
+ *
+ * THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS"
+ * AND ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+ * IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE
+ * ARE DISCLAIMED. IN NO EVENT SHALL THE COPYRIGHT OWNER OR CONTRIBUTORS BE
+ * LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR
+ * CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF
+ * SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR PROFITS; OR BUSINESS
+ * INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY OF LIABILITY, WHETHER IN
+ * CONTRACT, STRICT LIABILITY, OR TORT (INCLUDING NEGLIGENCE OR OTHERWISE)
+ * ARISING IN ANY WAY OUT OF THE USE OF THIS SOFTWARE, EVEN IF ADVISED OF THE
+ * POSSIBILITY OF SUCH DAMAGE.
+ */
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/NumericLimits.cuh>
+
+#include "type_shim.h"
+#include "compat.h"
+
+#define ALIGN_BYTES 16
+
+using Tensor = at::Tensor;
+using TensorList = at::TensorList;
+using ScalarType = at::ScalarType;
+using at::acc_type;
+
+template<typename T, typename AccumT, typename OutT>
+struct LogSoftMaxForwardEpilogue {
+  __device__ __forceinline__ LogSoftMaxForwardEpilogue(AccumT max_input, AccumT sum)
+    : logsum(max_input + std::log(sum)) {}
+
+  __device__ __forceinline__ LogSoftMaxForwardEpilogue(AccumT max_log_sum_exp)
+    : logsum(max_log_sum_exp) {}
+
+  __device__ __forceinline__ OutT operator()(T input) const {
+    return static_cast<OutT>(input - logsum);
+  }
+
+  const AccumT logsum;
+};
+
+template<typename T, typename AccumT, typename OutT>
+struct LogSoftMaxBackwardEpilogue {
+  __device__ __forceinline__ LogSoftMaxBackwardEpilogue(AccumT sum)
+    : sum(sum) {}
+
+  __device__ __forceinline__ T operator()(OutT gradOutput, OutT output) const {
+    return static_cast<T>(gradOutput - std::exp(static_cast<AccumT>(output)) * sum);
+  }
+
+  const AccumT sum;
+};
+
+
+
+const int max_threads = 1024;
+
+inline dim3 SoftMax_getBlockSize(int ILP, uint64_t dim_size) {
+  uint64_t block_size = 1;
+  uint64_t max_block_size = std::min(dim_size / ILP, static_cast<uint64_t>(max_threads));
+  while (block_size < (max_block_size/2)) block_size *= 2;
+  // Launch at least a single warp - the kernel assumes that.
+  block_size = std::max(block_size, static_cast<uint64_t>(32));
+  return dim3(block_size);
+}
+
+template<typename T>
+struct Add {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a + b;
+  }
+};
+
+template<typename T>
+struct Max {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a < b ? b : a;
+  }
+};
+
+
+////////////////////////////////////////////////////////////////////////////////
+// Regular kernel (fast when dim_size is large; requires inner_size == 1)
+////////////////////////////////////////////////////////////////////////////////
+
+
+template <typename T, typename AccumT>
+struct MaxFloat
+{
+  __device__ __forceinline__ AccumT operator()(AccumT max, T v) const {
+    return ::max(max, (AccumT)v);
+  }
+};
+
+template<typename T, typename AccumT>
+struct AddFloat
+{
+  __device__ __forceinline__ AccumT operator()(AccumT sum, T v) const {
+    return sum + v;
+  }
+};
+
+template<typename T, typename AccumT>
+struct SumExpFloat
+{
+  __device__ __forceinline__ SumExpFloat(AccumT v)
+    : max_k(v) {}
+
+  __device__ __forceinline__ AccumT operator()(AccumT sum, T v) const {
+    return sum + std::exp(v - max_k);
+  }
+
+  const AccumT max_k;
+};
+
+template <template<typename> class Reduction, typename AccumT>
+__device__ __forceinline__ AccumT
+blockReduce(AccumT* smem, AccumT val,
+            const Reduction<AccumT>& r,
+            AccumT defaultVal)
+{
+  // To avoid RaW races from chaining blockReduce calls together, we need a sync here
+  __syncthreads();
+
+  smem[threadIdx.x] = val;
+
+  __syncthreads();
+
+  AccumT warpVal = defaultVal;
+
+  // First warp will perform per-warp reductions for the remaining warps
+  uint32_t mask = (((uint64_t)1) << (blockDim.x / 32)) - 1;
+  if (threadIdx.x < 32) {
+    int lane = threadIdx.x % 32;
+    if (lane < blockDim.x / 32) {
+#pragma unroll
+      for (int i = 0; i < 32; ++i) {
+        warpVal = r(warpVal, smem[lane * 32 + i]);
+      }
+      __syncwarp(mask);
+      smem[lane] = warpVal;
+    }
+  }
+
+  __syncthreads();
+
+  // First thread will perform a reduction of the above per-warp reductions
+  AccumT blockVal = defaultVal;
+
+  if (threadIdx.x == 0) {
+    for (int i = 0; i < blockDim.x / 32; ++i) {
+      blockVal = r(blockVal, smem[i]);
+    }
+    smem[0] = blockVal;
+  }
+
+  // Sync and broadcast
+  __syncthreads();
+  return smem[0];
+}
+
+template <template<typename> class Reduction1, template<typename> class Reduction2, typename AccumT>
+__device__ __forceinline__ void
+blockReduce(AccumT* smem,
+            AccumT* reducVal1,
+            AccumT val1,
+            const Reduction1<AccumT>& r1,
+            AccumT defaultVal1,
+            AccumT* reducVal2,
+            AccumT val2,
+            const Reduction2<AccumT>& r2,
+            AccumT defaultVal2)
+{
+  // To avoid RaW races from chaining blockReduce calls together, we need a sync here
+  __syncthreads();
+
+  smem[threadIdx.x] = val1;
+  smem[blockDim.x + threadIdx.x] = val2;
+
+  __syncthreads();
+
+  AccumT warpVal1 = defaultVal1;
+  AccumT warpVal2 = defaultVal2;
+
+  // First warp will perform per-warp reductions for the remaining warps
+  uint32_t mask = (((uint64_t)1) << (blockDim.x / 32)) - 1;
+  if (threadIdx.x < 32) {
+    int lane = threadIdx.x % 32;
+    if (lane < blockDim.x / 32) {
+#pragma unroll
+      for (int i = 0; i < 32; ++i) {
+        warpVal1 = r1(warpVal1, smem[lane * 32 + i]);
+        warpVal2 = r2(warpVal2, smem[lane * 32 + i + blockDim.x]);
+      }
+      __syncwarp(mask);
+      smem[lane] = warpVal1;
+      smem[lane + blockDim.x] = warpVal2;
+    }
+  }
+
+  __syncthreads();
+
+  // First thread will perform a reduction of the above per-warp reductions
+  AccumT blockVal1 = defaultVal1;
+  AccumT blockVal2 = defaultVal2;
+
+  if (threadIdx.x == 0) {
+    for (int i = 0; i < blockDim.x / 32; ++i) {
+      blockVal1 = r1(blockVal1, smem[i]);
+      blockVal2 = r2(blockVal2, smem[i + blockDim.x]);
+    }
+    smem[0] = blockVal1;
+    smem[blockDim.x] = blockVal2;
+  }
+
+  // Sync and broadcast
+  __syncthreads();
+  *reducVal1 = smem[0];
+  *reducVal2 = smem[blockDim.x];
+  __syncthreads();
+}
+
+template <template<typename, typename> class Reduction, int ILP, typename T, typename AccumT>
+__device__ __forceinline__ AccumT
+ilpReduce(int shift,
+          T* data,
+          int size,
+          const Reduction<T, AccumT>& r,
+          AccumT defaultVal)
+{
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LoadT;
+  AccumT threadVal = defaultVal;
+  int offset = threadIdx.x;
+
+  // shift and do 1
+  if(shift > 0){
+    data -= shift;
+    size += shift;
+    if(threadIdx.x >= shift){
+      threadVal = r(threadVal, data[offset]);
+    }
+    size -= blockDim.x;
+    data += blockDim.x;
+  }
+  int last = size % (ILP * blockDim.x);
+
+  T v[ILP];
+  LoadT* value = reinterpret_cast<LoadT*>(&v);
+
+  for (; offset * ILP < (size - last); offset += blockDim.x) {
+    *value = reinterpret_cast<LoadT*>(data)[offset];
+
+    for (int j = 0; j < ILP; ++j) {
+      threadVal = r(threadVal, v[j]);
+    }
+  }
+
+  offset = size - last + threadIdx.x;
+  // Epilogue
+  for (; offset < size; offset += blockDim.x)
+    threadVal = r(threadVal, data[offset]);
+
+  return threadVal;
+}
+
+template <template<typename, typename> class Reduction1, template<typename, typename> class Reduction2, int ILP, typename T, typename AccumT>
+__device__ __forceinline__ void
+ilpReduce(int shift,
+          T* data,
+          int size,
+          AccumT* reducVal1,
+          const Reduction1<T, AccumT>& r1,
+          AccumT defaultVal1,
+          AccumT* reducVal2,
+          const Reduction2<T, AccumT>& r2,
+          AccumT defaultVal2)
+{
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LoadT;
+
+  AccumT threadVal1 = defaultVal1;
+  AccumT threadVal2 = defaultVal2;
+  int offset = threadIdx.x;
+
+  // shift and do 1
+  if(shift > 0){
+    data -= shift;
+    size += shift;
+    if(threadIdx.x >= shift){
+      threadVal1 = r1(threadVal1, data[offset]);
+      threadVal2 = r2(threadVal2, data[offset]);
+    }
+    size -= blockDim.x;
+    data += blockDim.x;
+  }
+  int last = size % (ILP * blockDim.x);
+
+  T v[ILP];
+  LoadT* value = reinterpret_cast<LoadT*>(&v);
+
+  for (; offset * ILP < (size - last); offset += blockDim.x) {
+    *value = reinterpret_cast<LoadT*>(data)[offset];
+
+    for (int j = 0; j < ILP; ++j) {
+      threadVal1 = r1(threadVal1, v[j]);
+      threadVal2 = r2(threadVal2, v[j]);
+    }
+  }
+
+  offset = size - last + threadIdx.x;
+  // Epilogue
+  for (; offset < size; offset += blockDim.x) {
+    threadVal1 = r1(threadVal1, data[offset]);
+    threadVal2 = r2(threadVal2, data[offset]);
+  }
+
+  *reducVal1 = threadVal1;
+  *reducVal2 = threadVal2;
+}
+
+template <int ILP, typename scalar_t, typename accscalar_t, typename outscalar_t, template <typename, typename, typename> class Epilogue>
+__global__ void
+cunn_SoftMaxXEntropyForward(
+    accscalar_t *losses,
+    outscalar_t *max_log_sum_exp,
+    scalar_t *input,
+    int64_t *labels,
+    int64_t classes,
+    const float smoothing)
+{
+  extern __shared__ unsigned char smem[];
+  auto sdata = reinterpret_cast<accscalar_t*>(smem);
+  // forward pointers to batch[blockIdx.x]
+  // each block handles a sample in the mini-batch
+  input += blockIdx.x * classes;
+  //output += blockIdx.x * classes;
+  const int shift = ((uint64_t)input) % ALIGN_BYTES / sizeof(scalar_t);
+
+  int64_t label = labels[blockIdx.x];
+
+  // find the max and sum
+  accscalar_t threadMax, threadSum, max_k, sum_k;
+  ilpReduce<MaxFloat, AddFloat, ILP, scalar_t, accscalar_t>(
+    shift, input, classes,
+    &threadMax, MaxFloat<scalar_t, accscalar_t>(),
+    -at::numeric_limits<accscalar_t>::max(),
+    &threadSum, AddFloat<scalar_t, accscalar_t>(),
+    static_cast<accscalar_t>(0));
+
+  blockReduce<Max, Add, accscalar_t>(
+      sdata,
+      &max_k, threadMax, Max<accscalar_t>(),
+      -at::numeric_limits<accscalar_t>::max(),
+      &sum_k, threadSum, Add<accscalar_t>(),
+      static_cast<accscalar_t>(0));
+
+  accscalar_t threadExp = ilpReduce<SumExpFloat, ILP, scalar_t, accscalar_t>(shift, input, classes, SumExpFloat<scalar_t, accscalar_t>(max_k), static_cast<accscalar_t>(0));
+  accscalar_t sumAll = blockReduce<Add, accscalar_t>(
+      sdata, threadExp, Add<accscalar_t>(), static_cast<accscalar_t>(0));
+
+  Epilogue<scalar_t, accscalar_t, outscalar_t> epilogue(max_k, sumAll);
+
+  // calculate per element loss with label smoothing
+  // reserve max + log_sum_exp for bprop
+  if (threadIdx.x == 0) {
+    accscalar_t log_prob = epilogue(static_cast<accscalar_t>(input[label]));
+    losses[blockIdx.x] = (max_k + std::log(sumAll) - sum_k / classes) \
+      * smoothing - log_prob * (1 - smoothing);
+    max_log_sum_exp[blockIdx.x] = max_k + std::log(sumAll);
+  }
+}
+
+template <int ILP, typename scalar_t, typename accscalar_t, typename outscalar_t>
+__device__ __forceinline__ void
+apply(scalar_t *gradInput,
+      scalar_t *logits,
+      outscalar_t *max_log_sum_exp,
+      outscalar_t *gradOutput,
+      int64_t *labels,
+      const float smoothing,
+      int classes)
+{
+  accscalar_t smooth_positives = 1.0 - smoothing;
+  accscalar_t smooth_negatives = smoothing / classes;
+  accscalar_t tmpGradOutput = gradOutput[blockIdx.x];
+  int64_t label = labels[blockIdx.x];
+  accscalar_t coeff = max_log_sum_exp[blockIdx.x];
+
+  int offset = threadIdx.x;
+  int last = classes % (ILP * blockDim.x);
+
+  for (; offset < classes - last; offset += blockDim.x * ILP) {
+    accscalar_t tmpLogits[ILP];
+
+#pragma unroll
+    for (int j = 0; j < ILP; ++j) {
+      tmpLogits[j] = static_cast<accscalar_t>(logits[offset + j * blockDim.x]);
+    }
+
+#pragma unroll
+    for (int j = 0; j < ILP; ++j)
+      gradInput[offset + j * blockDim.x] = tmpGradOutput * (
+        std::exp(tmpLogits[j] - coeff) - static_cast<accscalar_t>(
+          (offset + j * blockDim.x == label) ? 1 : 0) *
+        smooth_positives - smooth_negatives);
+  }
+
+  for (; offset < classes; offset += blockDim.x)
+    gradInput[offset] = tmpGradOutput * (std::exp(
+        static_cast<accscalar_t>(logits[offset]) - coeff) -
+        static_cast<accscalar_t>((offset == label) ? 1 : 0) *
+        smooth_positives - smooth_negatives);
+}
+
+
+template <int ILP, typename scalar_t, typename accscalar_t, typename outscalar_t>
+__device__ __forceinline__ void
+aligned_apply(int shift,
+              scalar_t *gradInput,
+              scalar_t *logits,
+              outscalar_t *max_log_sum_exp,
+              outscalar_t *gradOutput,
+              int64_t *labels,
+              const float smoothing,
+              int classes)
+{
+  accscalar_t smooth_positives = 1.0 - smoothing;
+  accscalar_t smooth_negatives = smoothing / classes;
+  accscalar_t tmpGradOutput = gradOutput[blockIdx.x];
+  int64_t label = labels[blockIdx.x];
+  accscalar_t coeff = max_log_sum_exp[blockIdx.x];
+
+  int offset = threadIdx.x;
+
+  // shift and do 1
+  if(shift > 0){
+    logits -= shift;
+    gradInput -= shift;
+    classes += shift;
+    if(threadIdx.x >= shift){
+      gradInput[offset] = tmpGradOutput * (std::exp(
+        static_cast<accscalar_t>(logits[offset]) - coeff) -
+        static_cast<accscalar_t>(((offset - shift) == label) ? 1 : 0) *
+        smooth_positives - smooth_negatives);
+    }
+    classes -= blockDim.x;
+    gradInput += blockDim.x;
+    logits += blockDim.x;
+    shift -= blockDim.x;
+  }
+
+  int last = classes % (ILP * blockDim.x);
+
+  typedef typename std::aligned_storage<ILP*sizeof(scalar_t), ILP*alignof(scalar_t)>::type LoadT;
+  // input
+  scalar_t v[ILP];
+  LoadT* value = reinterpret_cast<LoadT*>(&v);
+  // output
+  scalar_t r[ILP];
+  LoadT* result = reinterpret_cast<LoadT*>(&r);
+
+  for (; offset * ILP < (classes - last); offset += blockDim.x) {
+    *value = reinterpret_cast<LoadT*>(logits)[offset];
+
+#pragma unroll
+    for (int j = 0; j < ILP; ++j) {
+      r[j] = tmpGradOutput * (std::exp(
+          static_cast<accscalar_t>(v[j]) - coeff) -
+          static_cast<accscalar_t>(((ILP * offset + j - shift) == label) ? 1 : 0) *
+          smooth_positives - smooth_negatives);
+    }
+    reinterpret_cast<LoadT*>(gradInput)[offset] = *result;
+  }
+
+  offset = classes - last + threadIdx.x;
+  for (; offset < classes; offset += blockDim.x)
+    gradInput[offset] = tmpGradOutput * (std::exp(
+        static_cast<accscalar_t>(logits[offset]) - coeff) -
+        static_cast<accscalar_t>(((offset - shift) == label) ? 1 : 0) *
+        smooth_positives - smooth_negatives);
+
+}
+
+template <int ILP, typename scalar_t, typename accscalar_t, typename outscalar_t, template<typename, typename, typename> class Epilogue>
+__global__ void
+cunn_SoftMaxXEntropyBackward(
+    scalar_t *gradInput,
+    scalar_t *logits,
+    outscalar_t *max_log_sum_exp,
+    outscalar_t *gradOutput,
+    int64_t *labels,
+    const float smoothing,
+    int classes)
+{
+  gradInput += blockIdx.x * classes;
+  logits += blockIdx.x * classes;
+
+  // Do vectorized load/store when input/output have same alignment
+  const int shift = ((uint64_t)logits) % ALIGN_BYTES / sizeof(scalar_t);
+  const int shift_ = ((uint64_t)gradInput) % ALIGN_BYTES / sizeof(scalar_t);
+  if (shift == shift_){
+    aligned_apply<ILP, scalar_t, accscalar_t, outscalar_t>(shift, gradInput, logits, max_log_sum_exp, gradOutput, labels, smoothing, classes);
+  }
+  else {
+    apply<ILP, scalar_t, accscalar_t, outscalar_t>(gradInput, logits, max_log_sum_exp, gradOutput, labels, smoothing, classes);
+  }
+
+}
+
+template<template<typename, typename, typename> class Epilogue>
+std::vector<Tensor> host_softmax_xentropy(
+        const Tensor & input_,
+        const Tensor & labels_,
+        const float smoothing,
+        const bool half_to_float){
+  if (half_to_float) AT_ASSERTM(input_.type().scalarType() == ScalarType::Half,"conversion is supported for Half type only");
+  AT_ASSERTM(labels_.type().scalarType() == ScalarType::Long,"Label type should be CUDA Long");
+
+  auto input = input_.contiguous();
+  Tensor max_log_sum_exp = at::empty_like(labels_, half_to_float ? input.options().dtype(ScalarType::Float) : input.options());
+  Tensor losses = at::empty_like(labels_, input_.options().dtype(ScalarType::Float));
+
+  static_assert(std::is_same<acc_type<at::Half, true>, float>::value ||
+    std::is_same<acc_type<at::Half, true>, double>::value,
+    "accscalar_t for half should be float or double");
+  AT_ASSERTM(input.dim() == 2, "Currently only 2 dim input supported");
+  AT_ASSERTM(labels_.dim() == 1, "Labels should be 1 dimensional");
+  AT_ASSERTM(input.size(0) == labels_.size(0), "Input and label should have same number of examples");
+  AT_ASSERTM(input.numel() > 0, "Number of classes in input should not be 0");
+
+  const int64_t dim = 1;
+  int64_t outer_size = 1;
+  int64_t dim_size = input.size(dim);
+  int64_t inner_size = 1;
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  for (int64_t i = 0; i < dim; ++i)
+    outer_size *= input.size(i);
+  for (int64_t i = dim + 1; i < input.dim(); ++i)
+    inner_size *= input.size(i);
+  // This kernel spawns a block per each element in the batch.
+  // XXX: it assumes that inner_size == 1
+  TORCH_CHECK(inner_size == 1, "Currently only inner size 1 supported");
+
+  dim3 grid(outer_size);
+
+  using namespace at;
+  DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "host_softmax_xentropy",
+    using accscalar_t = at::acc_type<scalar_t_0, true>;
+    const int ILP = sizeof(float4)/sizeof(scalar_t_0);
+    dim3 block = SoftMax_getBlockSize(ILP, dim_size);
+    if (!half_to_float) {
+      cunn_SoftMaxXEntropyForward<ILP, scalar_t_0, accscalar_t, scalar_t_0, Epilogue>
+        <<<grid, block, 2 * block.x * sizeof(accscalar_t), stream>>>(
+          losses.DATA_PTR<accscalar_t>(), max_log_sum_exp.DATA_PTR<scalar_t_0>(),
+          input.DATA_PTR<scalar_t_0>(), labels_.DATA_PTR<int64_t>(),
+          dim_size, smoothing
+      );
+    } else {
+      cunn_SoftMaxXEntropyForward<ILP, scalar_t_0, accscalar_t, accscalar_t, Epilogue>
+        <<<grid, block, 2 * block.x * sizeof(accscalar_t), stream>>>(
+          losses.DATA_PTR<accscalar_t>(), max_log_sum_exp.DATA_PTR<accscalar_t>(),
+          input.DATA_PTR<scalar_t_0>(), labels_.DATA_PTR<int64_t>(),
+          dim_size, smoothing
+      );
+    }
+  );
+
+  C10_CUDA_CHECK(cudaGetLastError());
+
+  std::vector<at::Tensor> ret = {losses, max_log_sum_exp};
+  return ret;
+}
+
+template<template<typename, typename, typename> class Epilogue>
+Tensor host_softmax_xentropy_backward(
+    const at::Tensor &grad_loss,
+    const at::Tensor &logits_,
+    const at::Tensor &max_log_sum_exp,
+    const at::Tensor &labels,
+    const float smoothing,
+    bool half_to_float) {
+  const int64_t dim = 1;
+  Tensor gI = at::empty_like(logits_);
+  if (grad_loss.numel() == 0) {
+    return gI;
+  }
+
+  auto grad = grad_loss.contiguous();
+  auto logits = logits_.contiguous();
+
+  static_assert(std::is_same<acc_type<at::Half, true>, float>::value ||
+    std::is_same<acc_type<at::Half, true>, double>::value,
+    "accscalar_t for half should be float or double");
+  if (grad.dim() == 0) grad = grad.view(1);
+
+  AT_ASSERTM(logits_.dim() == 2, "Currently only 2 dim input supported");
+  AT_ASSERTM(labels.dim() == 1, "Labels should be 1 dimensional");
+  AT_ASSERTM(logits_.numel() > 0, "Number of classes in input should not be 0");
+  AT_ASSERTM(logits_.size(0) == labels.size(0), "Input and label should have same number of examples");
+  AT_ASSERTM(labels.size(0) == grad.size(0), "Label and loss should have same number of examples");
+
+  int64_t outer_size = 1;
+  int64_t dim_size = logits.size(dim);
+  int64_t inner_size = 1;
+  for (int64_t i = 0; i < dim; ++i)
+    outer_size *= logits.size(i);
+  for (int64_t i = dim + 1; i < logits.dim(); ++i)
+    inner_size *= logits.size(i);
+  // See descriptions of kernels above.
+  cudaStream_t stream = at::cuda::getCurrentCUDAStream();
+  TORCH_CHECK(inner_size == 1, "Currently only inner size 1 supported");
+
+  dim3 grid(outer_size);
+
+  DISPATCH_FLOAT_AND_HALF(gI.scalar_type(), 0, "host_softmax_xentropy_backward",
+    using accscalar_t = acc_type<scalar_t_0, true>;
+    const int ILP = sizeof(float4)/sizeof(scalar_t_0);
+    dim3 block = SoftMax_getBlockSize(ILP, dim_size);
+    if (!half_to_float) {
+      cunn_SoftMaxXEntropyBackward<ILP, scalar_t_0, accscalar_t, scalar_t_0, Epilogue>
+       <<<grid, block, block.x * sizeof(accscalar_t), stream>>>(
+          gI.DATA_PTR<scalar_t_0>(), logits.DATA_PTR<scalar_t_0>(),
+          max_log_sum_exp.DATA_PTR<scalar_t_0>(),
+          grad.DATA_PTR<scalar_t_0>(), labels.DATA_PTR<int64_t>(),
+          smoothing, dim_size
+      );
+    } else {
+      cunn_SoftMaxXEntropyBackward<ILP, scalar_t_0, accscalar_t, accscalar_t, Epilogue>
+       <<<grid, block, block.x * sizeof(accscalar_t), stream>>>(
+          gI.DATA_PTR<scalar_t_0>(), logits.DATA_PTR<scalar_t_0>(),
+          max_log_sum_exp.DATA_PTR<accscalar_t>(),
+          grad.DATA_PTR<accscalar_t>(), labels.DATA_PTR<int64_t>(),
+          smoothing, dim_size
+      );
+    }
+  );
+
+  C10_CUDA_CHECK(cudaGetLastError());
+  return gI;
+}
+
+std::vector<Tensor> softmax_xentropy_cuda(const Tensor &input, const Tensor &labels, const float smoothing, const bool half_to_float){
+  return host_softmax_xentropy<LogSoftMaxForwardEpilogue>(input, labels, smoothing, half_to_float);
+}
+
+at::Tensor softmax_xentropy_backward_cuda(
+    const at::Tensor &grad_loss,
+    const at::Tensor &logits,
+    const at::Tensor &max_log_sum_exp,
+    const at::Tensor &labels,
+    const float smoothing) {
+  bool half_to_float = grad_loss.type().scalarType() != logits.type().scalarType();
+  if (half_to_float) {
+     AT_ASSERTM((grad_loss.type().scalarType() == ScalarType::Float && logits.type().scalarType() == ScalarType::Half), "expected input and grad types to match, or input to be at::Half and grad to be at::Float");
+  }
+  return host_softmax_xentropy_backward<LogSoftMaxBackwardEpilogue>(grad_loss, logits, max_log_sum_exp, labels, smoothing, half_to_float);
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/examples/multihead_attn/func_test_multihead_attn.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/examples/multihead_attn/func_test_multihead_attn.py
new file mode 100644
index 0000000000..10407b9973
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/examples/multihead_attn/func_test_multihead_attn.py
@@ -0,0 +1,108 @@
+import torch
+import torch.nn.functional as F
+import argparse
+
+from apex.contrib.multihead_attn import SelfMultiheadAttn
+from apex.contrib.multihead_attn import EncdecMultiheadAttn
+
+parser = argparse.ArgumentParser(description='Multihead Attention Standalone Test')
+parser.add_argument('--seq-length', default=64, type=int, help='Sequence Length of Input')
+parser.add_argument('--num-seqs-start', default=5, type=int, help='Start Range of Number of Sequences')
+parser.add_argument('--num-seqs-stop', default=80, type=int, help='Stop Range of Number of Sequences')
+parser.add_argument('--num-seqs-inc', default=5, type=int, help='Range Increment of Number of Sequences')
+parser.add_argument('--trials', default=20, type=int, help='Number of Trials to Execute')
+parser.add_argument('--warmup-trials', default=5, type=int, help='Warmup Trials to discard')
+parser.add_argument('--layers', default=18, type=int, help='Attention Layers to Execute to Gain CPU/GPU Time Overlap')
+parser.add_argument('--seed-start', default=1, type=int, help='Attention Layers to Execute to Gain CPU/GPU Time Overlap')
+parser.add_argument('--seed-end', default=100, type=int, help='Attention Layers to Execute to Gain CPU/GPU Time Overlap')
+parser.add_argument('--hidden-dim', default=1024, type=int, help='Multihead Attention hidden dimension')
+parser.add_argument('--heads', default=16, type=int, help='Number of Multihead Attention heads')
+parser.add_argument('--encdec-attn', action='store_true', help='Use Encoder-Decoder Attention instead of Self Attention.')
+parser.add_argument('--norm-add', action='store_true', help='Include Layer Norm and Dropout-Add in Multihead Attention block.')
+parser.add_argument('--ref', action='store_true', help='Reference implementation in python pytorch.')
+parser.add_argument('--native', action='store_true', help='torch.nn.MultitheadAttention Version.')
+parser.add_argument('--fwd', action='store_true', help='Only execute Fwd Pass.')
+parser.add_argument('--eval', action='store_true', help='Inference only, no backward pass.')
+
+args = parser.parse_args()
+assert args.seq_length % 64 == 0, "Sequence Length should be a multiple of 64!"
+
+if not torch.cuda.is_available():
+    raise NotImplementedError('Running on CPU is not supported')
+torch.cuda.set_device(0)
+
+dropout_prob = 0.1
+
+for seed in range(args.seed_start, args.seed_end+1) :
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    ref_layer = None
+    if args.encdec_attn :
+        ref_layer = EncdecMultiheadAttn(args.hidden_dim, args.heads, dropout=dropout_prob, bias=False, include_norm_add=args.norm_add, impl='default')
+    else :
+        ref_layer = SelfMultiheadAttn(args.hidden_dim, args.heads, dropout=dropout_prob, bias=False, include_norm_add=args.norm_add, impl='default')
+    ref_layer.cuda()
+    ref_layer.half()
+    ref_layer.reset_parameters()
+
+    ref_inputs    = torch.randn(args.seq_length, args.num_seqs_start, args.hidden_dim, dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+    ref_inputs_kv = None
+    if args.encdec_attn :
+        ref_inputs_kv    = torch.randn(args.seq_length, args.num_seqs_start, args.hidden_dim, dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+
+    ref_grads         = torch.randn_like(ref_inputs)
+
+    ref_outputs,_ = ref_layer.forward(ref_inputs,
+                                      ref_inputs_kv,
+                                      ref_inputs_kv,
+                                      key_padding_mask=None,
+                                      need_weights=False,
+                                      attn_mask=None,
+                                      is_training=(not args.eval))
+
+    ref_outputs.backward(ref_grads)
+
+    torch.manual_seed(seed)
+    if torch.cuda.is_available():
+        torch.cuda.manual_seed_all(seed)
+    tst_layer = None
+    if args.encdec_attn :
+        tst_layer = EncdecMultiheadAttn(args.hidden_dim, args.heads, dropout=dropout_prob, bias=False, include_norm_add=args.norm_add, impl='fast')
+    else:
+        tst_layer = SelfMultiheadAttn(args.hidden_dim, args.heads, dropout=dropout_prob, bias=False, include_norm_add=args.norm_add, impl='fast')
+    tst_layer.cuda()
+    tst_layer.half()
+    tst_layer.reset_parameters()
+
+    tst_inputs    = torch.randn(args.seq_length, args.num_seqs_start, args.hidden_dim, dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+    tst_inputs_kv = None
+    if args.encdec_attn :
+        tst_inputs_kv    = torch.randn(args.seq_length, args.num_seqs_start, args.hidden_dim, dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+
+    assert torch.equal(ref_inputs,tst_inputs), "ERROR: Inputs are different!"
+
+    tst_grads         = torch.randn_like(tst_inputs)
+
+    tst_outputs,_ = tst_layer.forward(tst_inputs,
+                                      tst_inputs_kv,
+                                      tst_inputs_kv,
+                                      key_padding_mask=None,
+                                      need_weights=False,
+                                      attn_mask=None,
+                                      is_training=(not args.eval))
+
+    tst_outputs.backward(tst_grads)
+
+    fwd_close = torch.equal(ref_outputs, tst_outputs)
+    bwd_close = torch.equal(ref_inputs.grad, tst_inputs.grad)
+
+    diff_fwd = ref_outputs - tst_outputs
+    diff_cnt_fwd = diff_fwd.ne(0.0).sum()
+    diff_accum_fwd = diff_fwd.abs().sum()
+
+    diff_bwd = ref_inputs.grad - tst_inputs.grad
+    diff_cnt_bwd = diff_bwd.ne(0.0).sum()
+    diff_accum_bwd = diff_bwd.abs().sum()
+
+    print(">>> Seed: ", seed, fwd_close, diff_cnt_fwd.item(), diff_accum_fwd.item(), bwd_close, diff_cnt_bwd.item(), diff_accum_bwd.item())
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/examples/multihead_attn/perf_test_multihead_attn.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/examples/multihead_attn/perf_test_multihead_attn.py
new file mode 100644
index 0000000000..f81522ab37
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/examples/multihead_attn/perf_test_multihead_attn.py
@@ -0,0 +1,115 @@
+import torch
+import torch.nn.functional as F
+import argparse
+
+from apex.contrib.multihead_attn import SelfMultiheadAttn
+from apex.contrib.multihead_attn import EncdecMultiheadAttn
+
+parser = argparse.ArgumentParser(description='Multihead Attention Standalone Test')
+parser.add_argument('--seq-length', default=64, type=int, help='Sequence Length of Input')
+parser.add_argument('--num-seqs-start', default=10, type=int, help='Start Range of Number of Sequences')
+parser.add_argument('--num-seqs-stop', default=120, type=int, help='Stop Range of Number of Sequences')
+parser.add_argument('--num-seqs-inc', default=5, type=int, help='Range Increment of Number of Sequences')
+parser.add_argument('--trials', default=20, type=int, help='Number of Trials to Execute')
+parser.add_argument('--warmup-trials', default=5, type=int, help='Warmup Trials to discard')
+parser.add_argument('--layers', default=18, type=int, help='Attention Layers to Execute to Gain CPU/GPU Time Overlap')
+parser.add_argument('--hidden-dim', default=1024, type=int, help='Multihead Attention hidden dimension')
+parser.add_argument('--heads', default=16, type=int, help='Number of Multihead Attention heads')
+parser.add_argument('--encdec-attn', action='store_true', help='Use Encoder-Decoder Attention instead of Self Attention.')
+parser.add_argument('--norm-add', action='store_true', help='Include Layer Norm and Dropout-Add in Multihead Attention block.')
+parser.add_argument('--ref', action='store_true', help='Reference implementation in python pytorch.')
+parser.add_argument('--native', action='store_true', help='torch.nn.MultitheadAttention Version.')
+parser.add_argument('--fwd', action='store_true', help='Only execute Fwd Pass.')
+parser.add_argument('--biases', action='store_true', help='Execute multihead attention with Linear Biases.')
+
+args = parser.parse_args()
+
+if not torch.cuda.is_available():
+    raise NotImplementedError('Running on CPU is not supported')
+torch.cuda.set_device(0)
+
+torch.manual_seed(111)
+if torch.cuda.is_available():
+    torch.cuda.manual_seed_all(111)
+
+attn_layers = []
+for idx in range(0, args.layers) :
+    if args.encdec_attn :
+        if args.ref :
+            attn_layers.append(EncdecMultiheadAttn(args.hidden_dim, args.heads, dropout=0.1, bias=args.biases, include_norm_add=False, impl='default'))
+        else :
+            attn_layers.append(EncdecMultiheadAttn(args.hidden_dim, args.heads, dropout=0.1, bias=args.biases, include_norm_add=args.norm_add, impl='fast'))
+    else :
+        if args.native :
+            attn_layers.append(torch.nn.MultiheadAttention(args.hidden_dim, args.heads, dropout=0.1, bias=args.biases))
+        elif args.ref :
+            attn_layers.append(SelfMultiheadAttn(args.hidden_dim, args.heads, dropout=0.1, bias=args.biases, include_norm_add=args.norm_add, impl='default'))
+        else :
+            attn_layers.append(SelfMultiheadAttn(args.hidden_dim, args.heads, dropout=0.1, bias=args.biases, include_norm_add=args.norm_add, impl='fast'))
+    attn_layers[idx].cuda()
+    attn_layers[idx].half()
+    if not args.native :
+        attn_layers[idx].reset_parameters()
+
+start_evt_fwd = []
+start_evt_bwd = []
+stop_evt_bwd  = []
+for recorded_trial in range(0, args.trials) :
+    start_evt_fwd.append(torch.cuda.Event(enable_timing=True))
+    start_evt_bwd.append(torch.cuda.Event(enable_timing=True))
+    stop_evt_bwd.append(torch.cuda.Event(enable_timing=True))
+
+for sequences in range(args.num_seqs_start, args.num_seqs_stop + args.num_seqs_inc, args.num_seqs_inc) :
+    inputs        = torch.randn(args.seq_length, sequences, args.hidden_dim, dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+    grads         = torch.randn_like(inputs)
+   
+    for trial in range(0, args.trials + args.warmup_trials) :
+        layer_inputs  = inputs
+        evt_idx       = trial - args.warmup_trials
+    
+        if evt_idx >= 0 :
+            start_evt_fwd[evt_idx].record()
+    
+        for lyr_idx in range(0, args.layers) :
+            if args.native :
+                outputs,_ = attn_layers[lyr_idx].forward(layer_inputs, 
+                                                         layer_inputs, 
+                                                         layer_inputs, 
+                                                         key_padding_mask=None, 
+                                                         need_weights=False, 
+                                                         attn_mask=None)
+            else :
+                outputs,_ = attn_layers[lyr_idx].forward(layer_inputs, 
+                                                         layer_inputs, 
+                                                         layer_inputs,
+                                                         key_padding_mask=None, 
+                                                         need_weights=False, 
+                                                         attn_mask=None,
+                                                         is_training=True)
+            layer_inputs = outputs
+    
+        if evt_idx >= 0 :
+            start_evt_bwd[evt_idx].record()
+
+        if not args.fwd :
+            layer_inputs.backward(grads)
+    
+        if evt_idx >= 0 :
+            stop_evt_bwd[evt_idx].record()
+   
+    torch.cuda.synchronize()
+    elapsed_time_fwd = 0.0
+    elapsed_time_bwd = 0.0
+    for evt_idx in range(0, args.trials) :
+        elapsed_time_fwd += start_evt_fwd[evt_idx].elapsed_time(start_evt_bwd[evt_idx])
+        elapsed_time_bwd += start_evt_bwd[evt_idx].elapsed_time(stop_evt_bwd[evt_idx])
+   
+    print("[ {} Attn {} ]Total Tokens: {:4d} Sequences: {:3d} Sequence Length: {:3d} Fwd Time / Layer: {:.3f} ms Bwd Time / Layer: {:.3f} ms".format(
+        'Encdec' if args.encdec_attn else 'Self',              \
+        'Norm&Add' if args.norm_add else '',                   \
+        sequences*args.seq_length,                             \
+        sequences,                                             \
+        args.seq_length,                                       \
+        elapsed_time_fwd / ( args.trials * args.layers ),      \
+        elapsed_time_bwd / ( args.trials * args.layers )))
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/fmha/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/fmha/__init__.py
new file mode 100644
index 0000000000..ec2e9c66c8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/fmha/__init__.py
@@ -0,0 +1 @@
+from .fmha import FMHAFun
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/fmha/fmha.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/fmha/fmha.py
new file mode 100644
index 0000000000..57215a6fd6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/fmha/fmha.py
@@ -0,0 +1,74 @@
+###############################################################################
+# Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+# 
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#     * Redistributions of source code must retain the above copyright
+#       notice, this list of conditions and the following disclaimer.
+#     * Redistributions in binary form must reproduce the above copyright
+#       notice, this list of conditions and the following disclaimer in the
+#       documentation and/or other materials provided with the distribution.
+#     * Neither the name of the NVIDIA CORPORATION nor the
+#       names of its contributors may be used to endorse or promote products
+#       derived from this software without specific prior written permission.
+# 
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+###############################################################################
+
+
+import torch
+import torch.nn.functional as F
+import fmhalib as mha
+
+class FMHAFun(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, qkv, cu_seqlens, p_dropout, max_s, is_training):
+        batch_size = cu_seqlens.numel() - 1
+        if batch_size < 4:
+            context, S_dmask = mha.fwd_nl(qkv, cu_seqlens, p_dropout, max_s, is_training, None)
+        else:
+            context, S_dmask = mha.fwd(qkv, cu_seqlens, p_dropout, max_s, is_training, None)
+        ctx.save_for_backward(qkv, S_dmask)
+        ctx.cu_seqlens = cu_seqlens
+        ctx.p_dropout = p_dropout
+        ctx.max_s = max_s
+        return context
+    
+    @staticmethod
+    def backward(ctx, dout):
+        qkv, S_dmask = ctx.saved_tensors
+        batch_size = ctx.cu_seqlens.numel() - 1
+        if batch_size < 4:
+            dqkv, dp, _ = mha.bwd_nl(dout, qkv, S_dmask, ctx.cu_seqlens, ctx.p_dropout, ctx.max_s)
+        else:
+            dqkv, dp = mha.bwd(dout, qkv, S_dmask, ctx.cu_seqlens, ctx.p_dropout, ctx.max_s)
+
+        return dqkv, None, None, None, None, None, None
+
+class FMHA(torch.nn.Module):
+
+    def __init__(self, config):
+
+        super(FMHA, self).__init__()
+
+        self.p_dropout = config.attention_probs_dropout_prob
+        self.h = config.num_attention_heads
+        self.hidden_size = config.hidden_size
+        self.d = self.hidden_size // self.h
+        assert self.d * self.h == self.hidden_size, "Invalid hidden size/num_heads"
+
+    def forward(self, qkv, cu_seqlens, max_s, is_training=True):
+
+        ctx = FMHAFun.apply(qkv.view(-1, 3, self.h, self.d), cu_seqlens, self.p_dropout, max_s, is_training)
+
+        return ctx.view(-1, self.hidden_size)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/groupbn/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/groupbn/__init__.py
new file mode 100644
index 0000000000..2f85770668
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/groupbn/__init__.py
@@ -0,0 +1,9 @@
+try:
+    import torch
+    import bnp
+    from .batch_norm import BatchNorm2d_NHWC
+    del torch
+    del bnp
+    del batch_norm
+except ImportError as err:
+    print("apex was installed without --bnp flag, contrib.groupbn is not available")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/groupbn/batch_norm.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/groupbn/batch_norm.py
new file mode 100644
index 0000000000..17ef196b99
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/groupbn/batch_norm.py
@@ -0,0 +1,225 @@
+import torch
+import numpy as np
+from torch.nn.modules.batchnorm import _BatchNorm
+
+import bnp
+
+class bn_NHWC_impl(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, s, b, rm, riv, mini_m, mini_riv, ret_cta, mom, epsilon, fuse_relu, is_train, bn_group, my_data, pair_data, magic, pair_data2, pair_data3, fwd_occup, fwd_grid_x, bwd_occup, bwd_grid_x, multi_stream):
+        if is_train:
+            ctx.save_for_backward(x, s, b, rm, riv, mini_m, mini_riv)
+            ctx.epsilon = epsilon
+            ctx.momentum = mom
+            ctx.ret_cta = ret_cta
+            ctx.fuse_relu = fuse_relu
+            ctx.my_data = my_data
+            ctx.pair_data = pair_data
+            ctx.magic = magic
+            ctx.pair_data2 = pair_data2
+            ctx.pair_data3 = pair_data3
+            ctx.bn_group = bn_group
+            ctx.bwd_occup = bwd_occup
+            ctx.bwd_grid_x = bwd_grid_x
+            ctx.multi_stream = multi_stream
+
+            res =  bnp.bn_fwd_nhwc(x, s, b, rm, riv, mini_m, mini_riv, ret_cta, mom, epsilon, fuse_relu, my_data, pair_data, pair_data2, pair_data3, bn_group, magic, fwd_occup, fwd_grid_x, multi_stream)
+            return res
+        else:
+            return bnp.bn_fwd_eval_nhwc(x, s, b, rm, riv, ret_cta, bn_group, mom, epsilon, fuse_relu)
+
+    @staticmethod
+    def backward(ctx, grad_y):
+        x, s, b, rm, riv, mini_m, mini_riv = ctx.saved_variables
+        epsilon = ctx.epsilon
+        mom = ctx.momentum
+        ret_cta = ctx.ret_cta
+        fuse_relu = ctx.fuse_relu
+        my_data = ctx.my_data
+        pair_data = ctx.pair_data
+        magic = ctx.magic
+        pair_data2 = ctx.pair_data2
+        pair_data3 = ctx.pair_data3
+        bn_group = ctx.bn_group
+        bwd_occup = ctx.bwd_occup
+        bwd_grid_x = ctx.bwd_grid_x
+        multi_stream = ctx.multi_stream
+
+        dx, dscale, dbias = bnp.bn_bwd_nhwc(x, grad_y, s, b, rm, riv, mini_m, mini_riv, ret_cta, mom, epsilon, fuse_relu, my_data, pair_data, pair_data2, pair_data3, bn_group, magic, bwd_occup, bwd_grid_x, multi_stream)
+
+        return dx, dscale, dbias, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None
+
+
+class bn_addrelu_NHWC_impl(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, z, s, b, rm, riv, mini_m, mini_riv, grid_dim_y, ret_cta, mom, epsilon, is_train, bn_group, my_data, pair_data, magic, pair_data2, pair_data3, fwd_occup, fwd_grid_x, bwd_occup, bwd_grid_x, multi_stream):
+        if is_train:
+            bitmask = torch.cuda.IntTensor(((x.numel()+31)//32) * 2 * grid_dim_y)
+            ctx.save_for_backward(x, s, b, rm, riv, mini_m, mini_riv, bitmask)
+            ctx.epsilon = epsilon
+            ctx.momentum = mom
+            ctx.ret_cta = ret_cta
+            ctx.my_data = my_data
+            ctx.pair_data = pair_data
+            ctx.magic = magic
+            ctx.pair_data2 = pair_data2
+            ctx.pair_data3 = pair_data3
+            ctx.bn_group = bn_group
+            ctx.bwd_occup = bwd_occup
+            ctx.bwd_grid_x = bwd_grid_x
+            ctx.multi_stream = multi_stream
+
+            res =  bnp.bn_addrelu_fwd_nhwc(x, z, s, b, rm, riv, mini_m, mini_riv, bitmask, ret_cta, mom, epsilon, my_data, pair_data, pair_data2, pair_data3, bn_group, magic, fwd_occup, fwd_grid_x, multi_stream)
+            return res
+        else:
+            return bnp.bn_addrelu_fwd_eval_nhwc(x, z, s, b, rm, riv, ret_cta, bn_group, mom, epsilon)
+
+    @staticmethod
+    def backward(ctx, grad_y):
+        x, s, b, rm, riv, mini_m, mini_riv, bitmask = ctx.saved_variables
+        epsilon = ctx.epsilon
+        mom = ctx.momentum
+        ret_cta = ctx.ret_cta
+        my_data = ctx.my_data
+        pair_data = ctx.pair_data
+        magic = ctx.magic
+        pair_data2 = ctx.pair_data2
+        pair_data3 = ctx.pair_data3
+        bn_group = ctx.bn_group
+        bwd_occup = ctx.bwd_occup
+        bwd_grid_x = ctx.bwd_grid_x
+        multi_stream = ctx.multi_stream
+
+        dx, dz, dscale, dbias = bnp.bn_addrelu_bwd_nhwc(x, grad_y, s, b, rm, riv, mini_m, mini_riv, bitmask, ret_cta, mom, epsilon, my_data, pair_data, pair_data2, pair_data3, bn_group, magic, bwd_occup, bwd_grid_x, multi_stream)
+
+        return dx, dz, dscale, dbias, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None, None
+
+
+
+
+
+class BatchNorm2d_NHWC(_BatchNorm):
+    # if using BatchNorm2d_NHWC simultaneously with multiple streams set multi_stream to True
+    def __init__(self, num_features, fuse_relu=False, bn_group=1, max_cta_per_sm=2, cta_launch_margin=12, multi_stream=False):
+        super(BatchNorm2d_NHWC, self).__init__(num_features)
+
+        self.fuse_relu = fuse_relu
+        self.multi_stream = multi_stream
+
+        self.minibatch_mean = torch.cuda.FloatTensor(num_features)
+        self.minibatch_riv = torch.cuda.FloatTensor(num_features)
+
+        #defaut to distributed bn disabled
+        self.bn_group = bn_group
+        self.max_cta_per_sm = max_cta_per_sm        #used only in training fwd and bwd
+        self.cta_launch_margin = cta_launch_margin  #used only in training fwd and bwd
+        self.my_data = None
+        self.pair_data = None
+        self.pair_data2 = None
+        self.pair_data3 = None
+        self.local_rank = 0
+        self.magic = torch.IntTensor([0])
+
+        #calculate cta per sm occupancies
+        assert(max_cta_per_sm>0) # won't be able to do much with 0 CTAs :)
+        self.fwd_occupancy =  min(bnp.bn_fwd_nhwc_occupancy(), max_cta_per_sm)
+        self.bwd_occupancy =  min(bnp.bn_bwd_nhwc_occupancy(), max_cta_per_sm)
+        self.addrelu_fwd_occupancy =  min(bnp.bn_addrelu_fwd_nhwc_occupancy(), max_cta_per_sm)
+        self.addrelu_bwd_occupancy =  min(bnp.bn_addrelu_bwd_nhwc_occupancy(), max_cta_per_sm)
+
+        #calculate grid dimentions based on occupancy numbers
+        mp_count = torch.cuda.get_device_properties(None).multi_processor_count
+        self.fwd_grid_dim_x = max(mp_count*self.fwd_occupancy - cta_launch_margin , 1)
+        self.bwd_grid_dim_x = max(mp_count*self.bwd_occupancy - cta_launch_margin , 1)
+        self.addrelu_fwd_grid_dim_x = max(mp_count*self.addrelu_fwd_occupancy - cta_launch_margin , 1)
+        self.addrelu_bwd_grid_dim_x = max(mp_count*self.addrelu_bwd_occupancy - cta_launch_margin , 1)
+        self.grid_dim_y = (num_features + 63) // 64
+
+        # allocate scratch space used by implementation
+        # TODO: scratch space that is not supposed to be exposed at user code. We only need one time initialization, the
+        # same buffer could be reused in future iterations. Currently we exposed it here instead of requesting new
+        # buffer from cache allocator to avoid unnecessary initialization at future iterations.
+        self.ret_cta = torch.cuda.ByteTensor(8192).fill_(0)
+
+        #FIXME: turn pair handles into an array
+        if bn_group>1:
+            local_rank = torch.distributed.get_rank()
+            world_size = torch.distributed.get_world_size()          
+            assert(world_size >= bn_group)
+            assert(world_size % bn_group == 0)
+         
+            bn_sync_steps = 1
+            if (bn_group==4):
+                bn_sync_steps = 2
+            if (bn_group==8):
+                bn_sync_steps = 3
+
+            self.ipc_buffer = torch.cuda.ByteTensor(bnp.get_buffer_size(bn_sync_steps))
+            self.my_data = bnp.get_data_ptr(self.ipc_buffer)
+            # we are walking on very thin ice here by utilizing internal `_share_cuda_()`
+            self.storage = self.ipc_buffer.storage()
+            self.share_cuda = self.storage._share_cuda_()
+            internal_cuda_mem = self.share_cuda
+            # internal_cuda_mem[1]: ipc_mem_handle
+            my_handle = torch.cuda.ByteTensor(np.frombuffer(internal_cuda_mem[1], dtype=np.uint8))
+            # internal_cuda_mem[3]: offset
+            my_offset = torch.cuda.IntTensor([internal_cuda_mem[3]])
+
+            handles_all = torch.empty(world_size, my_handle.size(0), dtype=my_handle.dtype, device=my_handle.device)
+            handles_l = list(handles_all.unbind(0))
+            torch.distributed.all_gather(handles_l, my_handle)
+
+            offsets_all = torch.empty(world_size, my_offset.size(0), dtype=my_offset.dtype, device=my_offset.device)
+            offsets_l = list(offsets_all.unbind(0))
+            torch.distributed.all_gather(offsets_l, my_offset)
+
+            #whom do I actually care about? that would be local_rank XOR 1
+            self.pair_handle = handles_l[local_rank ^ 1].cpu().contiguous()
+            pair_offset = offsets_l[local_rank ^ 1].cpu()
+            self.pair_data = bnp.get_remote_data_ptr(self.pair_handle, pair_offset)
+
+            if bn_group>2:
+                self.pair_handle2 = handles_l[local_rank ^ 2].cpu().contiguous()
+                pair_offset2 = offsets_l[local_rank ^ 2].cpu()
+                self.pair_data2 = bnp.get_remote_data_ptr(self.pair_handle2, pair_offset2)
+
+            if bn_group>4:
+                self.pair_handle3 = handles_l[local_rank ^ 4].cpu().contiguous()
+                pair_offset3 = offsets_l[local_rank ^ 4].cpu()
+                self.pair_data3 = bnp.get_remote_data_ptr(self.pair_handle3, pair_offset3)
+
+            #FIXME: get magic value into C code and eliminate from here
+            self.magic = torch.IntTensor([2])
+            self.local_rank = local_rank
+
+
+    def forward(self, x, z=None):
+        if z is not None:
+            assert(self.fuse_relu==True)
+            return bn_addrelu_NHWC_impl.apply(x, z,
+                                  self.weight, self.bias,
+                                  self.running_mean, self.running_var,
+                                  self.minibatch_mean, self.minibatch_riv, self.grid_dim_y, self.ret_cta,
+                                  self.momentum,
+                                  self.eps, self.training, self.bn_group, self.my_data, self.pair_data, (self.magic), self.pair_data2, self.pair_data3,
+                                  self.addrelu_fwd_occupancy, self.addrelu_fwd_grid_dim_x,
+                                  self.addrelu_bwd_occupancy, self.addrelu_bwd_grid_dim_x,
+                                  self.multi_stream)
+        else:
+            return bn_NHWC_impl.apply(x,
+                                  self.weight, self.bias,
+                                  self.running_mean, self.running_var,
+                                  self.minibatch_mean, self.minibatch_riv, self.ret_cta,
+                                  self.momentum,
+                                  self.eps, self.fuse_relu, self.training, self.bn_group, self.my_data, self.pair_data, (self.magic), self.pair_data2, self.pair_data3,
+                                  self.fwd_occupancy, self.fwd_grid_dim_x,
+                                  self.bwd_occupancy, self.bwd_grid_dim_x,
+                                  self.multi_stream)
+
+    def __del__(self):
+        if self.bn_group>1:
+          bnp.close_remote_data(self.pair_handle)
+          if self.bn_group>2:
+              bnp.close_remote_data(self.pair_handle2)
+              if self.bn_group>4:
+                 bnp.close_remote_data(self.pair_handle3)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/layer_norm/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/layer_norm/__init__.py
new file mode 100644
index 0000000000..4bbc4763aa
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/layer_norm/__init__.py
@@ -0,0 +1 @@
+from .layer_norm import FastLayerNorm
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/layer_norm/layer_norm.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/layer_norm/layer_norm.py
new file mode 100644
index 0000000000..8a8d26d43b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/layer_norm/layer_norm.py
@@ -0,0 +1,53 @@
+import torch
+from torch.nn import init
+
+from apex._autocast_utils import _cast_if_autocast_enabled
+import fast_layer_norm
+
+
+class FastLayerNormFN(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, gamma, beta, epsilon):
+        x = x.contiguous()
+        gamma = gamma.contiguous()
+        beta = beta.contiguous()
+        hidden_size = gamma.numel()
+        xmat = x.view((-1, hidden_size))
+        ymat, mu, rsigma = fast_layer_norm.ln_fwd(xmat, gamma, beta, epsilon)
+        ctx.save_for_backward(x, gamma, mu, rsigma)
+        return ymat.view(x.shape)
+
+    @staticmethod
+    def backward(ctx, dy):
+        # assert dy.is_contiguous()
+        dy = dy.contiguous()  # this happens!
+        x, gamma, mu, rsigma = ctx.saved_tensors
+
+        hidden_size = gamma.numel()
+        xmat = x.view((-1, hidden_size))
+        dymat = dy.view(xmat.shape)
+        dxmat, dgamma, dbeta, _, _ = fast_layer_norm.ln_bwd(dymat, xmat, mu, rsigma, gamma)
+        dx = dxmat.view(x.shape)
+        return dx, dgamma, dbeta, None
+
+
+def _fast_layer_norm(x, weight, bias, epsilon):
+    args = _cast_if_autocast_enabled(x, weight, bias, epsilon)
+    with torch.cuda.amp.autocast(enabled=False):
+        return FastLayerNormFN.apply(*args)
+
+
+class FastLayerNorm(torch.nn.Module):
+    def __init__(self, hidden_size, eps=1e-5):
+        super().__init__()
+        self.epsilon = eps
+        self.weight = torch.nn.Parameter(torch.Tensor(hidden_size))
+        self.bias = torch.nn.Parameter(torch.Tensor(hidden_size))
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        init.ones_(self.weight)
+        init.zeros_(self.bias)
+
+    def forward(self, x):
+        return _fast_layer_norm(x, self.weight, self.bias, self.epsilon)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/README.md
new file mode 100644
index 0000000000..bf0f3a07fb
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/README.md
@@ -0,0 +1,60 @@
+# Fast Multihead Attention 
+
+This implementation has two main features :
+* A C++ implementation to avoid the CPU overheads of Pytorch found with smaller batch sizes.
+* The removal of all copies and transposes found in standard implementations of Multihead Attention.
+
+|                                            | Python Version | C++ Version |
+| :----------------------------------------- | :------------: | :---------: |
+| Layer Norm and Residual Add Variant        | X              | X           |
+| Includes Linear Biases                     | X              |             |
+| Reduces CPU Overheads                      |                | X           |
+| Fuses masking with Softmax                 |                | X           |
+| Removes Transposes and Copies              | X              | X           |
+| Includes Self and Encoder/Decoder Variants | X              | X           |
+
+## How to Instantiate
+
+`SelfMultiheadAttn(` _hidden dim_, _heads_, _dropout=prob_, _bias=bool_, _include_norm_add=bool_, _impl='fast'_ `)`
+`EncdecMultiheadAttn(` _hidden dim_, _heads_, _dropout=prob_, _bias=bool_, _include_norm_add=bool_, _impl='fast'_ `)`
+
+ `impl` has two options:
+ * `fast` uses C++ Version
+ * `default` uses Python Version
+
+## Instructions to build on Linux
+
+```
+$ git clone https://github.com/NVIDIA/apex
+$ cd apex
+$ pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--fast_multihead_attn" ./
+```
+## Try Performance Tests Yourself!
+Perf test script is found here!
+```
+cd contrib/examples/multihead_attn
+```
+#### Fast Multihead Attention
+```
+python perf_test_multihead_attn.py --ref
+```
+#### Fast Multihead Attention with C++ Implementation
+```
+python perf_test_multihead_attn.py
+```
+#### Compare with `torch.nn.MultiheadAttn`
+```
+python perf_test_multihead_attn.py --native
+```
+#### Test your own range!
+```
+python perf_test_multihead_attn.py --seq-length 64 --num-seqs-start 10 --num-seqs-stop 120 --num-seqs-inc 5
+```
+
+## Performance Comparisons
+
+* Performance was measured with 64 token sequence lengths on an NVIDIA TitanV card.
+* Time is measured across multiple layers to simulate an in model scenario.
+
+![Multihead Attention Forward](MHA_fwd.png)
+![Multihead Attention Backward](MHA_bwd.png)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/__init__.py
new file mode 100644
index 0000000000..7de3edc234
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/__init__.py
@@ -0,0 +1,3 @@
+from .self_multihead_attn import SelfMultiheadAttn
+from .encdec_multihead_attn import EncdecMultiheadAttn
+from .mask_softmax_dropout_func import fast_mask_softmax_dropout_func
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/encdec_multihead_attn.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/encdec_multihead_attn.py
new file mode 100644
index 0000000000..f8355664a2
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/encdec_multihead_attn.py
@@ -0,0 +1,141 @@
+import math
+
+import torch
+from torch import nn
+from torch.nn import Parameter
+import torch.nn.functional as F
+
+from .encdec_multihead_attn_func               import encdec_attn_func
+from .fast_encdec_multihead_attn_func          import fast_encdec_attn_func
+from .fast_encdec_multihead_attn_norm_add_func import fast_encdec_attn_norm_add_func
+from apex.normalization.fused_layer_norm       import FusedLayerNorm
+
+if hasattr(torch._C, '_jit_set_profiling_executor') :
+    torch._C._jit_set_profiling_executor(False)
+if hasattr(torch._C, '_jit_set_profiling_mode') :
+    torch._C._jit_set_profiling_mode(False)
+
+@torch.jit.script
+def jit_dropout_add(x, residual, prob, is_training):
+    # type: (Tensor, Tensor, float, bool) -> Tensor
+    out = F.dropout(x, p=prob, training=True)
+    out = residual + out
+    return out
+
+
+class EncdecMultiheadAttn(nn.Module):
+    """Multi-headed attention.
+
+    See "Attention Is All You Need" for more details.
+    """
+    def __init__(self, embed_dim, num_heads, dropout=0., bias=False, include_norm_add=False, impl='fast'):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+        self.bias = bias
+        self.include_norm_add = include_norm_add
+        self.impl = impl
+        self.scaling = self.head_dim**-0.5
+
+        self.in_proj_weight_q    = Parameter(torch.Tensor(embed_dim, embed_dim))
+        self.in_proj_weight_kv   = Parameter(torch.Tensor(2*embed_dim, embed_dim))
+        self.out_proj_weight     = Parameter(torch.Tensor(embed_dim, embed_dim))
+        if self.bias:
+            assert impl != 'fast', "ERROR! The Fast implementation does not support biases!"
+            self.in_proj_bias_q  = Parameter(torch.Tensor(embed_dim))
+            self.in_proj_bias_kv = Parameter(torch.Tensor(2*embed_dim))
+            self.out_proj_bias   = Parameter(torch.Tensor(embed_dim))
+        else:
+            self.register_parameter('in_proj_bias_q', None)
+            self.register_parameter('in_proj_bias_kv', None)
+            self.in_proj_bias_q  = None
+            self.in_proj_bias_kv = None
+            self.out_proj_bias   = None
+        if self.include_norm_add:
+            if impl == 'fast':
+                self.lyr_nrm_gamma_weights = Parameter(torch.Tensor(embed_dim))
+                self.lyr_nrm_beta_weights  = Parameter(torch.Tensor(embed_dim))
+                self.lyr_nrm               = None
+            else:
+                self.register_parameter('lyr_norm_gamma_weights', None)
+                self.register_parameter('lyr_norm_beta_weights', None)
+                self.lyr_nrm_gamma_weights = None
+                self.lyr_nrm_beta_weights  = None
+                self.lyr_nrm = FusedLayerNorm(embed_dim)
+        self.reset_parameters()
+
+        if self.include_norm_add:
+            if   impl == 'fast'    : self.attn_func = fast_encdec_attn_norm_add_func
+            elif impl == 'default' : self.attn_func = encdec_attn_func
+            else :                   assert False, "Unsupported impl: {} !".format(impl)
+        else:
+            if   impl == 'fast'    : self.attn_func = fast_encdec_attn_func
+            elif impl == 'default' : self.attn_func = encdec_attn_func
+            else :                   assert False, "Unsupported impl: {} !".format(impl)
+
+    def reset_parameters(self):
+        nn.init.xavier_uniform_(self.in_proj_weight_q)
+        # in_proj_weight_kv has shape [2 * hidden, hidden] but it should be
+        # initialized like a [hidden, hidden] matrix.
+        # sqrt(6 / (hidden + hidden)) / sqrt(6 / (2 * hidden + hidden)) = sqrt(1.5)
+        # therefore xavier_uniform gain should be set to sqrt(1.5).
+        nn.init.xavier_uniform_(self.in_proj_weight_kv, gain=math.sqrt(1.5))
+        nn.init.xavier_uniform_(self.out_proj_weight)
+        if self.bias:
+            nn.init.constant_(self.in_proj_bias_q, 0.)
+            nn.init.constant_(self.in_proj_bias_kv, 0.)
+            nn.init.constant_(self.out_proj_bias, 0.)
+        if self.include_norm_add:
+            if self.impl == 'fast' :
+                nn.init.ones_(self.lyr_nrm_gamma_weights)
+                nn.init.zeros_(self.lyr_nrm_beta_weights)
+            else:
+                self.lyr_nrm.reset_parameters()
+
+    def forward(self, query, key, value, key_padding_mask=None, need_weights=False, attn_mask=None, is_training=True):
+        """Input shape: Time x Batch x Channel
+
+        Self-attention can be implemented by passing in the same arguments for
+        query, key and value. Future timesteps can be masked with the
+        `mask_future_timesteps` argument. Padding elements can be excluded from
+        the key by passing a binary ByteTensor (`key_padding_mask`) with shape:
+        batch x src_len, where padding elements are indicated by 1s.
+        """
+
+        if key_padding_mask is not None:
+            assert (attn_mask is None), "ERROR attn_mask and key_padding_mask should not be both defined!"
+            mask = key_padding_mask
+        elif attn_mask is not None:
+            mask = attn_mask
+        else:
+            mask = None
+
+        if self.include_norm_add:
+            if self.impl == 'fast':
+                outputs = self.attn_func(attn_mask is not None, is_training, self.num_heads, query, key,
+                                         self.lyr_nrm_gamma_weights, self.lyr_nrm_beta_weights,
+                                         self.in_proj_weight_q, self.in_proj_weight_kv, self.out_proj_weight, mask, self.dropout)
+            else:
+                lyr_nrm_results = self.lyr_nrm(query)
+                outputs = self.attn_func(attn_mask is not None, is_training, self.num_heads, self.scaling, lyr_nrm_results, key,
+                                         self.in_proj_weight_q, self.in_proj_weight_kv, self.out_proj_weight,
+                                         self.in_proj_bias_q, self.in_proj_bias_kv, self.out_proj_bias,
+                                         mask, self.dropout)
+                if is_training:
+                    outputs = jit_dropout_add(outputs, query, self.dropout, is_training)
+                else:
+                    outputs = outputs + query
+        else:
+            if self.impl == 'fast':
+                outputs = self.attn_func(attn_mask is not None, is_training, self.num_heads, query, key,
+                                         self.in_proj_weight_q, self.in_proj_weight_kv, self.out_proj_weight, mask, self.dropout)
+            else:
+                outputs = self.attn_func(attn_mask is not None, is_training, self.num_heads, self.scaling, query, key,
+                                         self.in_proj_weight_q, self.in_proj_weight_kv, self.out_proj_weight,
+                                         self.in_proj_bias_q, self.in_proj_bias_kv, self.out_proj_bias,
+                                         mask, self.dropout)
+
+        return outputs,None
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/encdec_multihead_attn_func.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/encdec_multihead_attn_func.py
new file mode 100644
index 0000000000..5c5d9a4f3d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/encdec_multihead_attn_func.py
@@ -0,0 +1,268 @@
+import torch
+import torch.nn.functional as F
+
+
+class EncdecAttnFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, use_time_mask, is_training, heads, scale, inputs_q, inputs_kv,
+                input_weights_q, input_weights_kv, output_weights,
+                input_biases_q, input_biases_kv, output_biases,
+                mask, dropout_prob):
+        use_biases_t   = torch.tensor([input_biases_q is not None])
+        heads_t        = torch.tensor([heads])
+        scale_t        = torch.tensor([scale])
+        dropout_prob_t = torch.tensor([dropout_prob])
+        null_tensor    = torch.tensor([])
+        head_dim       = inputs_q.size(2) // heads
+
+        # Input Linear GEMM Q
+        # input1: (activations) [seql_q, seqs, embed_dim(1024)]
+        # input2: (weights)     [embed_dim (1024), embed_dim (1024)] (transpose [0,1])
+        # output:               [seql_q, seqs, embed_dim]
+        # GEMM: ( (seql_q*seqs) x embed_dim ) x ( embed_dim x embed_dim ) = (seql_q*seqs x embed_dim)
+        if use_biases_t[0]:
+            input_lin_q_results = torch.addmm(input_biases_q,
+                                              inputs_q.view(inputs_q.size(0) * inputs_q.size(1), inputs_q.size(2)),
+                                              input_weights_q.transpose(0,1),
+                                              beta=1., alpha=1.)
+        else:
+            input_lin_q_results = torch.mm(inputs_q.view(inputs_q.size(0) * inputs_q.size(1), inputs_q.size(2)), input_weights_q.transpose(0,1))
+        input_lin_q_results = input_lin_q_results.view(inputs_q.size(0), inputs_q.size(1), input_weights_q.size(0))
+        # Input Linear GEMM KV
+        # input1: (activations) [seql_k, seqs, embed_dim(1024)]
+        # input2: (weights)     [embed_dim*2 (2048), embed_dim (1024)] (transpose [0,1])
+        # output:               [seql_k, seqs, embed_dim*2]
+        # GEMM: ( (seql_k*seqs) x embed_dim ) x ( embed_dim x embed_dim*2 ) = (seql_k*seqs x embed_dim*2)
+        if use_biases_t[0]:
+            input_lin_kv_results = torch.addmm(input_biases_kv,
+                                               inputs_kv.view(inputs_kv.size(0) * inputs_kv.size(1), inputs_kv.size(2)),
+                                               input_weights_kv.transpose(0,1),
+                                               beta=1., alpha=1.)
+        else:
+            input_lin_kv_results = torch.mm(inputs_kv.view(inputs_kv.size(0) * inputs_kv.size(1), inputs_kv.size(2)), input_weights_kv.transpose(0,1))
+        input_lin_kv_results = input_lin_kv_results.view(inputs_kv.size(0), inputs_kv.size(1), input_weights_kv.size(0))
+
+        # Slice out k,v from one big Input Linear outuput (should only impact meta data, no copies!)
+        # Sequences and heads are combined to make the batch of the Batched GEMM
+        # input_lin_kv_results: [seql_k, seqs, heads(16), 2, head_dim(64)]
+        # input_lin_kv_results: [seql_k, batches=seqs*heads, 2, head_dim]
+        queries = input_lin_q_results.view(inputs_q.size(0), inputs_q.size(1)*heads, head_dim)
+        input_lin_kv_results = input_lin_kv_results.view(inputs_kv.size(0), inputs_kv.size(1)*heads, 2, head_dim)
+        keys    = input_lin_kv_results[:,:,0,:]
+        values  = input_lin_kv_results[:,:,1,:]
+
+        # Matmul1 Batched GEMMs
+        # The output tensor is specified prior to the Batch GEMM because baddbmm requires its specification
+        # baddbmm is used to apply the scale parameter via the Batched GEMM's alpha parameter instead of 
+        # a separate elementwise operation.
+        # Input1: (Queries) [seql_q, seqs*heads, head_dim] tranpose(0,1)
+        # Input2: (Keys)    [seql_k, seqs*heads, head_dim] transpose(0,1)
+        # output:           [seqs*heads, seql_q, seql_k]
+        # GEMM: Per batch: ( seql_q x head_dim ) x ( head_dim x seql_k ) = ( seql_q x seql_k )
+        matmul1_results = torch.empty((queries.size(1),queries.size(0),keys.size(0)), dtype=queries.dtype, device=torch.device('cuda'))
+        matmul1_results = torch.baddbmm(matmul1_results, queries.transpose(0,1), keys.transpose(0,1).transpose(1,2), out=matmul1_results, beta=0.0, alpha=scale_t[0])
+
+        if mask is not None:
+            # Self Attention Time Mask
+            if use_time_mask:
+                assert (len(mask.size()) == 2), "Timing mask is not 2D!"
+                assert (mask.size(0) == mask.size(1)), "Sequence length should match!"
+                mask = mask.to(torch.bool)
+                matmul1_results = matmul1_results.masked_fill_(mask, float('-inf'))
+            # Key Padding Mask
+            else:
+                batches,seql_q,seql_k = matmul1_results.size()
+                seqs = int(batches / heads)
+                matmul1_results = matmul1_results.view(seqs, heads, seql_q, seql_k)
+                mask = mask.to(torch.bool)
+                matmul1_results = matmul1_results.masked_fill_(mask.unsqueeze(1).unsqueeze(2), float('-inf'))
+                matmul1_results = matmul1_results.view(seqs*heads, seql_q, seql_k)
+
+        softmax_results = F.softmax(matmul1_results, dim=-1)
+
+        # Dropout - is not executed for inference
+        if is_training:
+            dropout_results,dropout_mask = torch._fused_dropout(softmax_results, p=(1.-dropout_prob_t[0]))
+        else:
+            dropout_results = softmax_results
+            dropout_mask    = null_tensor
+
+        # Matmul2 Batched GEMMs
+        # The output tensor specification is needed here to specify the non-standard output.
+        # Given that pytorch cannot currently perform autograd with an output tensor specified,
+        # this requires a backward pass specified.
+        # Input1: from_softmax [seqs*heads, seql_q, seql_k]
+        # Input2: (values)     [seql_v, seqs*heads, head_dim] transpose(0,1)
+        # Output:              [seql_q, seqs*heads, head_dim] transpose(0,1)
+        # GEMM: Per batch: ( seql_q x seql_k ) x ( seql_k x head_dim ) = (seql_q x head_dim)
+        matmul2_results = torch.empty((dropout_results.size(1), dropout_results.size(0), values.size(2)), dtype=dropout_results.dtype, device=torch.device('cuda')).transpose(1,0)
+        matmul2_results = torch.bmm(dropout_results, values.transpose(0,1), out=matmul2_results)
+        matmul2_results = matmul2_results.transpose(0, 1).contiguous().view(inputs_q.size(0), inputs_q.size(1), inputs_q.size(2))
+
+        # Output Linear GEMM
+        # Input1: (activations) [seql_q, seqs, embed_dim=heads*head_dim]
+        # Input2: (weights)     [ embed_dim, embed_dim ] transpose(0,1)
+        # Output:               [ seql_q, seqs, embed_dim ]
+        # GEMM: ( seql_q*seqs x embed_dim ) x ( embed_dim x embed_dim ) = ( seql_q*seqs x embed_dim )
+        if use_biases_t[0]:
+            outputs = torch.addmm(output_biases,
+                                  matmul2_results.view(inputs_q.size(0) * inputs_q.size(1), inputs_q.size(2)),
+                                  output_weights.transpose(0,1),
+                                  beta=1., alpha=1.)
+        else:
+            outputs = torch.mm(matmul2_results.view(inputs_q.size(0) * inputs_q.size(1), inputs_q.size(2)), output_weights.transpose(0,1))
+        outputs = outputs.view(inputs_q.size(0), inputs_q.size(1), output_weights.size(0))
+
+        ctx.save_for_backward(use_biases_t,                             \
+                              heads_t,                                  \
+                              scale_t,                                  \
+                              matmul2_results,                          \
+                              dropout_results,                          \
+                              softmax_results,                          \
+                              input_lin_q_results,                      \
+                              input_lin_kv_results,                     \
+                              inputs_q,                                 \
+                              inputs_kv,                                \
+                              input_weights_q,                          \
+                              input_weights_kv,                         \
+                              output_weights,                           \
+                              dropout_mask,                             \
+                              dropout_prob_t)
+
+        return outputs.detach()
+    
+    @staticmethod
+    def backward(ctx, output_grads):
+        use_biases_t,                                                   \
+        heads_t,                                                        \
+        scale_t,                                                        \
+        matmul2_results,                                                \
+        dropout_results,                                                \
+        softmax_results,                                                \
+        input_lin_q_results,                                            \
+        input_lin_kv_results,                                           \
+        inputs_q,                                                       \
+        inputs_kv,                                                      \
+        input_weights_q,                                                \
+        input_weights_kv,                                               \
+        output_weights,                                                 \
+        dropout_mask,                                                   \
+        dropout_prob_t          = ctx.saved_tensors
+
+        head_dim                = inputs_q.size(2) // heads_t[0]
+
+        # Slice out k,v from one big Input Linear outuput (should only impact meta data, no copies!)
+        # Sequences and heads are combined to make the batch of the Batched GEMM
+        # input_lin_kv_results: [seql_k, seqs, heads(16), 2, head_dim(64)]
+        # input_lin_kv_results: [seql_k, batches=seqs*heads, 2, head_dim]
+        queries = input_lin_q_results.view(inputs_q.size(0), inputs_q.size(1)*heads_t[0], head_dim)
+        input_lin_kv_results = input_lin_kv_results.view(inputs_kv.size(0), inputs_kv.size(1)*heads_t[0], 2, head_dim)
+        keys    = input_lin_kv_results[:,:,0,:]
+        values  = input_lin_kv_results[:,:,1,:]
+
+        # Slice out k,v from one big set of gradients entering the input linear's bprop  (should only impact meta data, no copies!)
+        # The gradients are identical in size to the Input Linear outputs.
+        # The tensor is declared before hand to properly slice out query, key, and value grads.
+        input_lin_kv_results_grads = torch.empty_like(input_lin_kv_results)
+        queries_grads              = torch.empty_like(queries)
+        keys_grads                 = input_lin_kv_results_grads[:,:,0,:]
+        values_grads               = input_lin_kv_results_grads[:,:,1,:]
+
+        # Output Linear GEMM - DGRAD
+        # Input1: (data grads)  [seql_q, seqs, embed_dim=heads*head_dim]
+        # Input2: (weights)     [ embed_dim, embed_dim ]
+        # Output:               [ seql_q, seqs, embed_dim ]
+        # GEMM: ( seql_q*seqs x embed_dim ) x ( embed_dim x embed_dim ) = ( seql_q*seqs x embed_dim )
+        output_lin_grads = torch.mm(output_grads.view(output_grads.size(0) * output_grads.size(1), output_grads.size(2)), output_weights)
+        output_lin_grads = output_lin_grads.view(output_grads.size(0), output_grads.size(1), output_weights.size(1))
+        # Output Linear GEMM - WGRAD
+        # Input1: (data grads)  [seql_q*seqs, embed_dim=heads*head_dim] transpose(0,1)
+        # Input2: (activations) [seql_q*seqs, embed_dim ]
+        # Output:               [ seql_q, seqs, embed_dim ]
+        # GEMM: ( embed_dim x seql_q*seqs ) x ( seql_q*seqs x embed_dim ) = ( embed_dim x embed_dim )
+        output_weight_grads = torch.mm(output_grads.view(output_grads.size(0) * output_grads.size(1), output_grads.size(2)).transpose(0,1),
+                                       matmul2_results.view(matmul2_results.size(0) * matmul2_results.size(1), matmul2_results.size(2)))
+        output_lin_grads = output_lin_grads.view(output_grads.size(0), output_grads.size(1)*heads_t[0], head_dim).transpose(0,1)
+
+        if use_biases_t[0]:
+            output_bias_grads = torch.sum(output_grads.view(output_grads.size(0) * output_grads.size(1), output_grads.size(2)), 0)
+        else:
+            output_bias_grads = None
+
+        # Matmul2 - DGRAD1
+        # Input1: (data grads)  [seql_q, seqs*heads, head_dim] transpose(0,1)
+        # Input2: (activations) [seql_k, seqs*heads, head_dim] transpose(0,1).transpose(1,2)
+        # Output:               [seqs*heads, seql_q, seql_k]
+        # GEMM: Per batch: ( seql_q x head_dim ) x ( head_dim x seql_k ) = ( seql_q x seql_k )
+        matmul2_dgrad1 = torch.bmm(output_lin_grads, values.transpose(0,1).transpose(1,2))
+        # Matmul2 - DGRAD2
+        # Input1: (data grads)  [seql_q, seqs*heads, head_dim] transpose(0,1)
+        # Input2: (activations) [seql_k, seqs*heads, head_dim] transpose(0,1).transpose(1,2)
+        # Output:               [seqs*heads, seql_q, seql_k]
+        # GEMM: Per batch: ( seql_q x head_dim ) x ( head_dim x seql_k ) = ( seql_q x seql_k )
+        values_grads   = torch.bmm(dropout_results.transpose(1,2), output_lin_grads, out=values_grads.transpose(0,1))
+
+        # Mask and Scaling for Dropout (not a publically documented op)
+        dropout_grads = torch._masked_scale(matmul2_dgrad1, dropout_mask, 1.0/(1.0-dropout_prob_t[0]))
+
+        # Softmax Grad (not a publically documented op)
+        softmax_grads = torch._softmax_backward_data(dropout_grads, softmax_results, -1, softmax_results)
+
+        # Matmul1 - DGRAD1
+        # Input1: (data grads)  [seqs*heads, seql_q, seql_k] 
+        # Input2: (activations) [seql_k, seqs*heads, head_dim] transpose(0,1)
+        # Output:               [seqs*heads, seql_q, head_dim] transpose(0,1)
+        # GEMM: Per batch: ( seql_q x seql_k ) x ( seql_k x head_dim ) = ( seql_q x head_dim )
+        queries_grads = torch.baddbmm(queries_grads.transpose(0,1), softmax_grads, keys.transpose(0,1),
+                                      out=queries_grads.transpose(0,1), beta=0.0, alpha=scale_t[0])
+        # Matmul1 - DGRAD2
+        # Input1: (data grads)  [seqs*heads, seql_q, seql_k] transpose(1,2)
+        # Input2: (activations) [seql_q, seqs*heads, head_dim] transpose(0,1)
+        # Output:               [seqs*heads, seql_k, head_dim] transpose(0,1)
+        # GEMM: Per batch: ( seql_k x seql_q ) x ( seql_q x head_dim ) = ( seql_k x head_dim )
+        keys_grads    = torch.baddbmm(keys_grads.transpose(0,1), softmax_grads.transpose(1,2), queries.transpose(0,1),
+                                      out=keys_grads.transpose(0,1), beta=0.0, alpha=scale_t[0])
+
+        # Input Q Linear GEMM - DGRAD
+        # input1: (data grads) [seql_q, seqs, embed_dim(1024)]
+        # input2: (weights)    [embed_dim (1024), embed_dim (1024)] 
+        # output:              [seql_q, seqs, embed_dim]
+        # GEMM: ( (seql_q*seqs) x embed_dim ) x ( embed_dim x embed_dim ) = (seql_q*seqs x embed_dim)
+        queries_grads  = queries_grads.transpose(0,1).view(inputs_q.size(0)*inputs_q.size(1), heads_t[0]*head_dim)
+        input_q_grads = torch.mm(queries_grads, input_weights_q)
+        input_q_grads = input_q_grads.view(inputs_q.size(0), inputs_q.size(1), inputs_q.size(2))
+        # Input KV Linear GEMM - DGRAD
+        # input1: (data grads) [seql_k, seqs, 2*embed_dim(2048)]
+        # input2: (weights)    [embed_dim*2 (2048), embed_dim (1024)] 
+        # output:              [seql_k, seqs, embed_dim]
+        # GEMM: ( (seql_k*seqs) x 2*embed_dim ) x ( 2*embed_dim x embed_dim ) = (seql_k*seqs x embed_dim)
+        input_lin_kv_results_grads = input_lin_kv_results_grads.view(inputs_kv.size(0)*inputs_kv.size(1), heads_t[0]*2*head_dim)
+        input_kv_grads = torch.mm(input_lin_kv_results_grads, input_weights_kv)
+        input_kv_grads = input_kv_grads.view(inputs_kv.size(0), inputs_kv.size(1), inputs_kv.size(2))
+        # Input Q Linear GEMM - WGRAD
+        # input1: (data grads)  [seql_q*seqs, embed_dim(1024)]
+        # input2: (activations) [seql_q*seqs, embed_dim(1024)] 
+        # output:               [embed_dim, embed_dim]
+        # GEMM: ( embed_dim x seql_q*seqs ) x ( seql_q*seqs x embed_dim ) = (embed_dim x embed_dim)
+        input_weight_q_grads = torch.mm(queries_grads.transpose(0,1), inputs_q.view(inputs_q.size(0)*inputs_q.size(1), inputs_q.size(2)))
+        # Input KV Linear GEMM - WGRAD
+        # input1: (data grads)  [seql_k*seqs, 2*embed_dim(2048)]
+        # input2: (activations) [seql_k*seqs, embed_dim(1024)] 
+        # output:               [2*embed_dim, embed_dim]
+        # GEMM: ( 2*embed_dim x seql_k*seqs ) x ( seql_k*seqs x embed_dim ) = (2*embed_dim x embed_dim)
+        input_weight_kv_grads = torch.mm(input_lin_kv_results_grads.transpose(0,1), inputs_kv.view(inputs_kv.size(0)*inputs_kv.size(1), inputs_kv.size(2)))
+
+        if use_biases_t[0]:
+            input_bias_grads_q = torch.sum(queries_grads, 0)
+            input_bias_grads_kv = torch.sum(input_lin_kv_results_grads, 0)
+        else:
+            input_bias_grads_q = None
+            input_bias_grads_kv = None
+
+        return None, None, None, None,                                            \
+               input_q_grads, input_kv_grads,                                     \
+               input_weight_q_grads, input_weight_kv_grads, output_weight_grads,  \
+               input_bias_grads_q, input_bias_grads_kv, output_bias_grads,        \
+               None, None
+
+encdec_attn_func = EncdecAttnFunc.apply
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_encdec_multihead_attn_func.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_encdec_multihead_attn_func.py
new file mode 100644
index 0000000000..d56837aace
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_encdec_multihead_attn_func.py
@@ -0,0 +1,88 @@
+import torch
+import fast_encdec_multihead_attn
+
+
+class FastEncdecAttnFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, use_time_mask, is_training, heads, inputs_q, inputs_kv, input_weights_q, input_weights_kv, output_weights, pad_mask, dropout_prob):
+        heads_t        = torch.tensor([heads])
+        dropout_prob_t = torch.tensor([dropout_prob])
+        null_tensor    = torch.tensor([])
+        use_mask       = (pad_mask is not None)
+
+        input_lin_q_results,                                            \
+        input_lin_kv_results,                                           \
+        softmax_results,                                                \
+        dropout_results,                                                \
+        dropout_mask,                                                   \
+        matmul2_results,                                                \
+        outputs =                                                       \
+            fast_encdec_multihead_attn.forward(                         \
+                              use_mask,                                 \
+                              use_time_mask,                            \
+                              is_training,                              \
+                              heads,                                    \
+                              inputs_q,                                 \
+                              inputs_kv,                                \
+                              input_weights_q,                          \
+                              input_weights_kv,                         \
+                              output_weights,                           \
+                              pad_mask if use_mask else null_tensor,    \
+                              dropout_prob)
+
+        ctx.save_for_backward(heads_t,                                  \
+                              matmul2_results,                          \
+                              dropout_results,                          \
+                              softmax_results,                          \
+                              input_lin_q_results,                      \
+                              input_lin_kv_results,                     \
+                              inputs_q,                                 \
+                              inputs_kv,                                \
+                              input_weights_q,                          \
+                              input_weights_kv,                         \
+                              output_weights,                           \
+                              dropout_mask,                             \
+                              dropout_prob_t)
+
+        return outputs.detach()
+
+    @staticmethod
+    def backward(ctx, output_grads):
+        heads_t,                                                        \
+        matmul2_results,                                                \
+        dropout_results,                                                \
+        softmax_results,                                                \
+        input_lin_q_results,                                            \
+        input_lin_kv_results,                                           \
+        inputs_q,                                                       \
+        inputs_kv,                                                      \
+        input_weights_q,                                                \
+        input_weights_kv,                                               \
+        output_weights,                                                 \
+        dropout_mask,                                                   \
+        dropout_prob_t      = ctx.saved_tensors
+
+        input_q_grads,                                                  \
+        input_kv_grads,                                                 \
+        input_weight_q_grads,                                           \
+        input_weight_kv_grads,                                          \
+        output_weight_grads =                                           \
+            fast_encdec_multihead_attn.backward(                        \
+                              heads_t[0],                               \
+                              output_grads,                             \
+                              matmul2_results,                          \
+                              dropout_results,                          \
+                              softmax_results,                          \
+                              input_lin_q_results,                      \
+                              input_lin_kv_results,                     \
+                              inputs_q,                                 \
+                              inputs_kv,                                \
+                              input_weights_q,                          \
+                              input_weights_kv,                         \
+                              output_weights,                           \
+                              dropout_mask,                             \
+                              dropout_prob_t[0])
+
+        return None, None, None, input_q_grads, input_kv_grads, input_weight_q_grads, input_weight_kv_grads, output_weight_grads, None, None
+
+fast_encdec_attn_func = FastEncdecAttnFunc.apply
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_encdec_multihead_attn_norm_add_func.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_encdec_multihead_attn_norm_add_func.py
new file mode 100644
index 0000000000..213102bf15
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_encdec_multihead_attn_norm_add_func.py
@@ -0,0 +1,130 @@
+# Copyright (c) 2017-present, Facebook, Inc.
+# All rights reserved.
+#
+# This source code is licensed under the license found in the LICENSE file in
+# the root directory of this source tree. An additional grant of patent rights
+# can be found in the PATENTS file in the same directory.
+
+import torch
+import fast_encdec_multihead_attn_norm_add
+
+
+class FastEncdecAttnNormAddFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, use_time_mask, is_training, heads, inputs_q, inputs_kv, lyr_nrm_gamma_weights, lyr_nrm_beta_weights, input_weights_q, input_weights_kv, output_weights, pad_mask, dropout_prob):
+        heads_t        = torch.tensor([heads])
+        dropout_prob_t = torch.tensor([dropout_prob])
+        null_tensor    = torch.tensor([])
+        use_mask       = (pad_mask is not None)
+
+        lyr_nrm_results,                                                \
+        lyr_nrm_mean,                                                   \
+        lyr_nrm_invvar,                                                 \
+        input_lin_q_results,                                            \
+        input_lin_kv_results,                                           \
+        softmax_results,                                                \
+        dropout_results,                                                \
+        dropout_mask,                                                   \
+        matmul2_results,                                                \
+        dropout_add_mask,                                               \
+        outputs =                                                       \
+            fast_encdec_multihead_attn_norm_add.forward(                \
+                              use_mask,                                 \
+                              use_time_mask,                            \
+                              is_training,                              \
+                              heads,                                    \
+                              inputs_q,                                 \
+                              inputs_kv,                                \
+                              lyr_nrm_gamma_weights,                    \
+                              lyr_nrm_beta_weights,                     \
+                              input_weights_q,                          \
+                              input_weights_kv,                         \
+                              output_weights,                           \
+                              pad_mask if use_mask else null_tensor,    \
+                              dropout_prob)
+
+        ctx.save_for_backward(heads_t,                                  \
+                              matmul2_results,                          \
+                              dropout_results,                          \
+                              softmax_results,                          \
+                              input_lin_q_results,                      \
+                              input_lin_kv_results,                     \
+                              lyr_nrm_results,                          \
+                              lyr_nrm_mean,                             \
+                              lyr_nrm_invvar,                           \
+                              inputs_q,                                 \
+                              inputs_kv,                                \
+                              lyr_nrm_gamma_weights,                    \
+                              lyr_nrm_beta_weights,                     \
+                              input_weights_q,                          \
+                              input_weights_kv,                         \
+                              output_weights,                           \
+                              dropout_mask,                             \
+                              dropout_add_mask,                         \
+                              dropout_prob_t)
+
+        return outputs.detach()
+
+    @staticmethod
+    def backward(ctx, output_grads):
+        heads_t,                                                        \
+        matmul2_results,                                                \
+        dropout_results,                                                \
+        softmax_results,                                                \
+        input_lin_q_results,                                            \
+        input_lin_kv_results,                                           \
+        lyr_nrm_results,                                                \
+        lyr_nrm_mean,                                                   \
+        lyr_nrm_invvar,                                                 \
+        inputs_q,                                                       \
+        inputs_kv,                                                      \
+        lyr_nrm_gamma_weights,                                          \
+        lyr_nrm_beta_weights,                                           \
+        input_weights_q,                                                \
+        input_weights_kv,                                               \
+        output_weights,                                                 \
+        dropout_mask,                                                   \
+        dropout_add_mask,                                               \
+        dropout_prob_t         = ctx.saved_tensors
+
+        input_q_grads,                                                  \
+        input_kv_grads,                                                 \
+        lyr_nrm_gamma_grads,                                            \
+        lyr_nrm_beta_grads,                                             \
+        input_weight_q_grads,                                           \
+        input_weight_kv_grads,                                          \
+        output_weight_grads    =                                        \
+            fast_encdec_multihead_attn_norm_add.backward(               \
+                              heads_t[0],                               \
+                              output_grads,                             \
+                              matmul2_results,                          \
+                              dropout_results,                          \
+                              softmax_results,                          \
+                              input_lin_q_results,                      \
+                              input_lin_kv_results,                     \
+                              lyr_nrm_results,                          \
+                              lyr_nrm_mean,                             \
+                              lyr_nrm_invvar,                           \
+                              inputs_q,                                 \
+                              inputs_kv,                                \
+                              lyr_nrm_gamma_weights,                    \
+                              lyr_nrm_beta_weights,                     \
+                              input_weights_q,                          \
+                              input_weights_kv,                         \
+                              output_weights,                           \
+                              dropout_mask,                             \
+                              dropout_add_mask,                         \
+                              dropout_prob_t[0])
+
+        #import pdb; pdb.set_trace()
+        return None, None, None,                                        \
+               input_q_grads,                                           \
+               input_kv_grads,                                          \
+               lyr_nrm_gamma_grads,                                     \
+               lyr_nrm_beta_grads,                                      \
+               input_weight_q_grads,                                    \
+               input_weight_kv_grads,                                   \
+               output_weight_grads,                                     \
+               None, None
+
+fast_encdec_attn_norm_add_func = FastEncdecAttnNormAddFunc.apply
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_self_multihead_attn_func.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_self_multihead_attn_func.py
new file mode 100644
index 0000000000..1c75f9ead1
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_self_multihead_attn_func.py
@@ -0,0 +1,196 @@
+import torch
+import fast_self_multihead_attn
+import fast_self_multihead_attn_bias
+import fast_self_multihead_attn_bias_additive_mask
+
+class FastSelfAttnFunc(torch.autograd.Function) :
+    @staticmethod
+    def forward(ctx, use_time_mask, is_training, heads, inputs, input_weights, output_weights, input_biases, output_biases, pad_mask, mask_additive, dropout_prob):
+        use_biases_t   = torch.tensor([input_biases is not None])
+        heads_t        = torch.tensor([heads])
+        dropout_prob_t = torch.tensor([dropout_prob])
+        null_tensor    = torch.tensor([])
+        use_mask       = (pad_mask is not None)
+        mask_additive_t= torch.tensor([mask_additive])
+
+        if use_biases_t[0]:
+            if not mask_additive:
+                input_lin_results,                                              \
+                softmax_results,                                                \
+                dropout_results,                                                \
+                dropout_mask,                                                   \
+                matmul2_results,                                                \
+                outputs =                                                       \
+                    fast_self_multihead_attn_bias.forward(                           \
+                                      use_mask,                                 \
+                                      use_time_mask,                            \
+                                      is_training,                              \
+                                      heads,                                    \
+                                      inputs,                                   \
+                                      input_weights,                            \
+                                      output_weights,                           \
+                                      input_biases,                           \
+                                      output_biases,                           \
+                                      pad_mask if use_mask else null_tensor,    \
+                                      dropout_prob)
+                ctx.save_for_backward(use_biases_t,                                  \
+                              heads_t,                          \
+                              matmul2_results,                          \
+                              dropout_results,                          \
+                              softmax_results,                          \
+                              null_tensor,                          \
+                              null_tensor,                          \
+                              mask_additive_t,                          \
+                              input_lin_results,                        \
+                              inputs,                                   \
+                              input_weights,                            \
+                              output_weights,                           \
+                              dropout_mask,                             \
+                              dropout_prob_t)
+
+            else:
+                input_lin_results,                                              \
+                bmm1_results,                                                \
+                dropout_results,                                                \
+                dropout_mask,                                                   \
+                matmul2_results,                                                \
+                outputs =                                                       \
+                    fast_self_multihead_attn_bias_additive_mask.forward(                           \
+                                      use_mask,                                 \
+                                      use_time_mask,                            \
+                                      is_training,                              \
+                                      heads,                                    \
+                                      inputs,                                   \
+                                      input_weights,                            \
+                                      output_weights,                           \
+                                      input_biases,                           \
+                                      output_biases,                           \
+                                      pad_mask if use_mask else null_tensor,    \
+                                      dropout_prob)
+                ctx.save_for_backward(use_biases_t,                                  \
+                                      heads_t,                          \
+                                      matmul2_results,                          \
+                                      dropout_results,                          \
+                                      null_tensor,                          \
+                                      bmm1_results,                          \
+                                      pad_mask,                          \
+                                      mask_additive_t,                          \
+                                      input_lin_results,                        \
+                                      inputs,                                   \
+                                      input_weights,                            \
+                                      output_weights,                           \
+                                      dropout_mask,                             \
+                                      dropout_prob_t)
+
+
+        else:
+            input_lin_results,                                              \
+            softmax_results,                                                \
+            dropout_results,                                                \
+            dropout_mask,                                                   \
+            matmul2_results,                                                \
+            outputs =                                                       \
+                fast_self_multihead_attn.forward(                           \
+                                  use_mask,                                 \
+                                  use_time_mask,                            \
+                                  is_training,                              \
+                                  heads,                                    \
+                                  inputs,                                   \
+                                  input_weights,                            \
+                                  output_weights,                           \
+                                  pad_mask if use_mask else null_tensor,    \
+                                  dropout_prob)
+            ctx.save_for_backward(use_biases_t,                                  \
+                          heads_t,                          \
+                          matmul2_results,                          \
+                          dropout_results,                          \
+                          softmax_results,                          \
+                          null_tensor,                          \
+                          null_tensor,                          \
+                          mask_additive_t,                          \
+                          input_lin_results,                        \
+                          inputs,                                   \
+                          input_weights,                            \
+                          output_weights,                           \
+                          dropout_mask,                             \
+                          dropout_prob_t)
+        return outputs.detach()
+
+    @staticmethod
+    def backward(ctx, output_grads):
+        use_biases_t,                                                        \
+        heads_t,                                                        \
+        matmul2_results,                                                \
+        dropout_results,                                                \
+        softmax_results,                                                \
+        bmm1_results,                                                \
+        pad_mask,                                                \
+        mask_additive_t,                                                \
+        input_lin_results,                                              \
+        inputs,                                                         \
+        input_weights,                                                  \
+        output_weights,                                                 \
+        dropout_mask,                                                   \
+        dropout_prob_t      = ctx.saved_tensors
+
+        if use_biases_t[0]:
+            if not mask_additive_t[0]:
+                input_grads,                                                    \
+                input_weight_grads,                                             \
+                output_weight_grads,                                           \
+                input_bias_grads,                                                   \
+                output_bias_grads =                                                    \
+                    fast_self_multihead_attn_bias.backward(                          \
+                                      heads_t[0],                               \
+                                      output_grads,                             \
+                                      matmul2_results,                          \
+                                      dropout_results,                          \
+                                      softmax_results,                          \
+                                      input_lin_results,                        \
+                                      inputs,                                   \
+                                      input_weights,                            \
+                                      output_weights,                           \
+                                      dropout_mask,                             \
+                                      dropout_prob_t[0])
+
+            else:
+                input_grads,                                                    \
+                input_weight_grads,                                             \
+                output_weight_grads,                                           \
+                input_bias_grads,                                                   \
+                output_bias_grads =                                                    \
+                    fast_self_multihead_attn_bias_additive_mask.backward(                          \
+                                      heads_t[0],                               \
+                                      output_grads,                             \
+                                      matmul2_results,                          \
+                                      dropout_results,                          \
+                                      bmm1_results,                          \
+                                      pad_mask,                          \
+                                      input_lin_results,                        \
+                                      inputs,                                   \
+                                      input_weights,                            \
+                                      output_weights,                           \
+                                      dropout_mask,                             \
+                                      dropout_prob_t[0])
+                    
+        else:
+            input_bias_grads = None                                                    
+            output_bias_grads = None
+            input_grads,                                                    \
+            input_weight_grads,                                             \
+            output_weight_grads =                                           \
+                fast_self_multihead_attn.backward(                          \
+                                  heads_t[0],                               \
+                                  output_grads,                             \
+                                  matmul2_results,                          \
+                                  dropout_results,                          \
+                                  softmax_results,                          \
+                                  input_lin_results,                        \
+                                  inputs,                                   \
+                                  input_weights,                            \
+                                  output_weights,                           \
+                                  dropout_mask,                             \
+                                  dropout_prob_t[0])
+        return None, None, None, input_grads, input_weight_grads, output_weight_grads,input_bias_grads, output_bias_grads, None, None, None
+
+fast_self_attn_func = FastSelfAttnFunc.apply
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_self_multihead_attn_norm_add_func.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_self_multihead_attn_norm_add_func.py
new file mode 100644
index 0000000000..9a37985cd9
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/fast_self_multihead_attn_norm_add_func.py
@@ -0,0 +1,106 @@
+import torch
+import fast_self_multihead_attn_norm_add
+
+
+class FastSelfAttnNormAddFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, use_time_mask, is_training, heads, inputs, lyr_nrm_gamma_weights, lyr_nrm_beta_weights, input_weights, output_weights, pad_mask, dropout_prob):
+        heads_t        = torch.tensor([heads])
+        dropout_prob_t = torch.tensor([dropout_prob])
+        null_tensor    = torch.tensor([])
+        use_mask       = (pad_mask is not None)
+
+        lyr_nrm_results,                                                \
+        lyr_nrm_mean,                                                   \
+        lyr_nrm_invvar,                                                 \
+        input_lin_results,                                              \
+        softmax_results,                                                \
+        dropout_results,                                                \
+        dropout_mask,                                                   \
+        matmul2_results,                                                \
+        dropout_add_mask,                                               \
+        outputs =                                                       \
+             fast_self_multihead_attn_norm_add.forward(                 \
+                              use_mask,                                 \
+                              use_time_mask,                            \
+                              is_training,                              \
+                              heads,                                    \
+                              inputs,                                   \
+                              lyr_nrm_gamma_weights,                    \
+                              lyr_nrm_beta_weights,                     \
+                              input_weights,                            \
+                              output_weights,                           \
+                              pad_mask if use_mask else null_tensor,    \
+                              dropout_prob)
+
+        ctx.save_for_backward(heads_t,                                  \
+                              matmul2_results,                          \
+                              dropout_results,                          \
+                              softmax_results,                          \
+                              input_lin_results,                        \
+                              lyr_nrm_results,                          \
+                              lyr_nrm_mean,                             \
+                              lyr_nrm_invvar,                           \
+                              inputs,                                   \
+                              lyr_nrm_gamma_weights,                    \
+                              lyr_nrm_beta_weights,                     \
+                              input_weights,                            \
+                              output_weights,                           \
+                              dropout_mask,                             \
+                              dropout_add_mask,                         \
+                              dropout_prob_t)
+
+        return outputs.detach()
+
+    @staticmethod
+    def backward(ctx, output_grads):
+        heads_t,                                                        \
+        matmul2_results,                                                \
+        dropout_results,                                                \
+        softmax_results,                                                \
+        input_lin_results,                                              \
+        lyr_nrm_results,                                                \
+        lyr_nrm_mean,                                                   \
+        lyr_nrm_invvar,                                                 \
+        inputs,                                                         \
+        lyr_nrm_gamma_weights,                                          \
+        lyr_nrm_beta_weights,                                           \
+        input_weights,                                                  \
+        output_weights,                                                 \
+        dropout_mask,                                                   \
+        dropout_add_mask,                                               \
+        dropout_prob_t          = ctx.saved_tensors
+
+        input_grads,                                                    \
+        lyr_nrm_gamma_grads,                                            \
+        lyr_nrm_beta_grads,                                             \
+        input_weight_grads,                                             \
+        output_weight_grads    =                                        \
+            fast_self_multihead_attn_norm_add.backward(                 \
+                              heads_t[0],                               \
+                              output_grads,                             \
+                              matmul2_results,                          \
+                              dropout_results,                          \
+                              softmax_results,                          \
+                              input_lin_results,                        \
+                              lyr_nrm_results,                          \
+                              lyr_nrm_mean,                             \
+                              lyr_nrm_invvar,                           \
+                              inputs,                                   \
+                              lyr_nrm_gamma_weights,                    \
+                              lyr_nrm_beta_weights,                     \
+                              input_weights,                            \
+                              output_weights,                           \
+                              dropout_mask,                             \
+                              dropout_add_mask,                         \
+                              dropout_prob_t[0])
+
+        return None, None, None,                                        \
+               input_grads,                                             \
+               lyr_nrm_gamma_grads,                                     \
+               lyr_nrm_beta_grads,                                      \
+               input_weight_grads,                                      \
+               output_weight_grads,                                     \
+               None, None
+
+fast_self_attn_norm_add_func = FastSelfAttnNormAddFunc.apply
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/mask_softmax_dropout_func.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/mask_softmax_dropout_func.py
new file mode 100644
index 0000000000..516e6d859d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/mask_softmax_dropout_func.py
@@ -0,0 +1,81 @@
+import torch
+import fast_mask_softmax_dropout
+import fast_additive_mask_softmax_dropout
+
+
+class MaskSoftmaxDropout(torch.autograd.Function) :
+    @staticmethod
+    def forward(ctx, is_training, heads, inputs, pad_mask, mask_additive, dropout_prob):
+        heads_t        = torch.tensor([heads])
+        dropout_prob_t = torch.tensor([dropout_prob])
+        null_tensor    = torch.tensor([])
+        use_mask       = (pad_mask is not None)
+        use_mask_t     = torch.tensor([use_mask])
+        mask_additive_t     = torch.tensor([mask_additive])
+
+        if mask_additive:
+            dropout_results,                                                \
+            dropout_mask,                                                   \
+            softmax_results =                                                \
+                    fast_additive_mask_softmax_dropout.forward(                           \
+                                      use_mask,                                 \
+                                      is_training,                              \
+                                      heads,                                    \
+                                      inputs,                                   \
+                                      pad_mask if use_mask else null_tensor,    \
+                                      dropout_prob)
+        else:
+            dropout_results,                                                \
+            dropout_mask,                                                   \
+            softmax_results =                                                \
+                    fast_mask_softmax_dropout.forward(                           \
+                                      use_mask,                                 \
+                                      is_training,                              \
+                                      heads,                                    \
+                                      inputs,                                   \
+                                      pad_mask if use_mask else null_tensor,    \
+                                      dropout_prob)
+        
+        ctx.save_for_backward(
+                              use_mask_t,                                    \
+                              heads_t,                                 \
+                              softmax_results,                          \
+                              dropout_mask,                             \
+                              pad_mask if use_mask else null_tensor,        \
+                              mask_additive_t,        \
+                              dropout_prob_t)
+
+        return dropout_results.detach()
+
+    @staticmethod
+    def backward(ctx, output_grads):
+        use_mask_t, \
+        heads_t,   \
+        softmax_results,                                                \
+        dropout_mask,                                              \
+        pad_mask,                                                   \
+        mask_additive_t,                                                   \
+        dropout_prob_t      = ctx.saved_tensors
+
+        if mask_additive_t[0]:
+            input_grads =                                                    \
+                fast_additive_mask_softmax_dropout.backward(                          \
+                                  use_mask_t[0],                             \
+                                  heads_t[0],                             \
+                                  output_grads,                             \
+                                  softmax_results,                          \
+                                  dropout_mask,                             \
+                                  dropout_prob_t[0])
+        else:
+            input_grads =                                                    \
+                fast_mask_softmax_dropout.backward(                          \
+                                  use_mask_t[0],                             \
+                                  heads_t[0],                             \
+                                  output_grads,                             \
+                                  softmax_results,                          \
+                                  dropout_mask,                             \
+                                  pad_mask,                             \
+                                  dropout_prob_t[0])
+        return None, None, input_grads, None, None, None
+
+fast_mask_softmax_dropout_func = MaskSoftmaxDropout.apply
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/self_multihead_attn.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/self_multihead_attn.py
new file mode 100644
index 0000000000..6b9714be07
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/self_multihead_attn.py
@@ -0,0 +1,178 @@
+import math
+
+import torch
+from torch import nn
+from torch.nn import Parameter
+import torch.nn.functional as F
+
+from .self_multihead_attn_func               import self_attn_func
+from .fast_self_multihead_attn_func          import fast_self_attn_func
+from .fast_self_multihead_attn_norm_add_func import fast_self_attn_norm_add_func
+from apex.normalization.fused_layer_norm     import FusedLayerNorm
+
+if hasattr(torch._C, '_jit_set_profiling_executor') :
+    torch._C._jit_set_profiling_executor(False)
+if hasattr(torch._C, '_jit_set_profiling_mode') :
+    torch._C._jit_set_profiling_mode(False)
+
+@torch.jit.script
+def jit_dropout_add(x, residual, prob, is_training):
+    # type: (Tensor, Tensor, float, bool) -> Tensor
+    out = F.dropout(x, p=prob, training=True)
+    out = residual + out
+    return out
+
+
+class SelfMultiheadAttn(nn.Module):
+    """Multi-headed attention.
+
+    See "Attention Is All You Need" for more details.
+    """
+    def __init__(self, embed_dim, num_heads, dropout=0., bias=False, include_norm_add=False, impl='fast', separate_qkv_params=False, mask_additive=False):
+        super().__init__()
+        self.embed_dim = embed_dim
+        self.num_heads = num_heads
+        self.dropout = dropout
+        self.head_dim = embed_dim // num_heads
+        assert self.head_dim * num_heads == self.embed_dim, "embed_dim must be divisible by num_heads"
+        self.bias = bias
+        self.include_norm_add = include_norm_add
+        self.impl = impl
+        self.scaling = self.head_dim**-0.5
+        self.separate_qkv_params = separate_qkv_params
+        self.mask_additive = mask_additive
+        if mask_additive:
+            assert self.include_norm_add == False, "additive mask not supported with layer norm"
+            assert impl == 'default' or (impl == 'fast' and bias), "additive mask not supported for fast mode without bias"
+        if separate_qkv_params:
+            self.q_weight  = Parameter(torch.Tensor(embed_dim, embed_dim))
+            self.k_weight  = Parameter(torch.Tensor(embed_dim, embed_dim))
+            self.v_weight  = Parameter(torch.Tensor(embed_dim, embed_dim))
+        else:
+            self.in_proj_weight  = Parameter(torch.Tensor(3*embed_dim, embed_dim))
+        self.out_proj_weight = Parameter(torch.Tensor(embed_dim, embed_dim))
+        if self.bias:
+            if separate_qkv_params:
+                self.q_bias  = Parameter(torch.Tensor(embed_dim))
+                self.k_bias  = Parameter(torch.Tensor(embed_dim))
+                self.v_bias  = Parameter(torch.Tensor(embed_dim))
+            else:
+                self.in_proj_bias = Parameter(torch.Tensor(3*embed_dim))
+            self.out_proj_bias = Parameter(torch.Tensor(embed_dim))
+        else:
+            if separate_qkv_params:
+                self.register_parameter('q_bias', None)
+                self.register_parameter('k_bias', None)
+                self.register_parameter('v_bias', None)
+                self.q_bias = None
+                self.k_bias = None
+                self.v_bias = None
+            else:
+                self.register_parameter('in_proj_bias', None)
+                self.in_proj_bias = None
+            self.register_parameter('out_proj_bias', None)
+            self.out_proj_bias = None
+        if self.include_norm_add:
+            if impl == 'fast':
+                self.lyr_nrm_gamma_weights = Parameter(torch.Tensor(embed_dim))
+                self.lyr_nrm_beta_weights  = Parameter(torch.Tensor(embed_dim))
+                self.lyr_nrm               = None
+            else:
+                self.register_parameter('lyr_norm_gamma_weights', None)
+                self.register_parameter('lyr_norm_beta_weights', None)
+                self.lyr_nrm_gamma_weights = None
+                self.lyr_nrm_beta_weights  = None
+                self.lyr_nrm = FusedLayerNorm(embed_dim)
+        self.reset_parameters()
+
+        if self.include_norm_add:
+            if   impl == 'fast'    : self.attn_func = fast_self_attn_norm_add_func
+            elif impl == 'default' : self.attn_func = self_attn_func
+            else :                   assert False, "Unsupported impl: {} !".format(impl)
+        else:
+            if   impl == 'fast'    : self.attn_func = fast_self_attn_func
+            elif impl == 'default' : self.attn_func = self_attn_func
+            else :                   assert False, "Unsupported impl: {} !".format(impl)
+
+    def reset_parameters(self):
+        if self.separate_qkv_params:
+            nn.init.xavier_uniform_(self.q_weight)
+            nn.init.xavier_uniform_(self.k_weight)
+            nn.init.xavier_uniform_(self.v_weight)
+        else:
+            # in_proj_weight has shape [3 * hidden, hidden] but it should be
+            # initialized like a [hidden, hidden] matrix.
+            # sqrt(6 / (hidden + hidden)) / sqrt(6 / (3 * hidden + hidden)) = sqrt(2)
+            # therefore xavier_uniform gain should be set to sqrt(2).
+            nn.init.xavier_uniform_(self.in_proj_weight, gain=math.sqrt(2))
+        nn.init.xavier_uniform_(self.out_proj_weight)
+        if self.bias:
+            if self.separate_qkv_params:
+                nn.init.constant_(self.q_bias, 0.)
+                nn.init.constant_(self.k_bias, 0.)
+                nn.init.constant_(self.v_bias, 0.)
+            else:
+                nn.init.constant_(self.in_proj_bias, 0.)
+            nn.init.constant_(self.out_proj_bias, 0.)
+        if self.include_norm_add:
+            if self.impl == 'fast':
+                nn.init.ones_(self.lyr_nrm_gamma_weights)
+                nn.init.zeros_(self.lyr_nrm_beta_weights)
+            else:
+                self.lyr_nrm.reset_parameters()
+
+    def forward(self, query, key, value, key_padding_mask=None, need_weights=False, attn_mask=None, is_training=True):
+        """Input shape: Time x Batch x Channel
+
+        Self-attention can be implemented by passing in the same arguments for
+        query, key and value. Future timesteps can be masked with the
+        `mask_future_timesteps` argument. Padding elements can be excluded from
+        the key by passing a binary ByteTensor (`key_padding_mask`) with shape:
+        batch x src_len, where padding elements are indicated by 1s.
+        """
+        if self.separate_qkv_params:
+            input_weights = torch.cat([self.q_weight.view(self.num_heads,1,self.head_dim,self.embed_dim), self.k_weight.view(self.num_heads,1,self.head_dim,self.embed_dim), self.v_weight.view(self.num_heads,1,self.head_dim,self.embed_dim)], dim=1).reshape(3*self.embed_dim,self.embed_dim).contiguous()
+        else: 
+            input_weights = self.in_proj_weight
+        if self.bias:
+            if self.separate_qkv_params:
+                input_bias = torch.cat([self.q_bias.view(self.num_heads,1,self.head_dim), self.k_bias.view(self.num_heads,1,self.head_dim), self.v_bias.view(self.num_heads,1,self.head_dim)],dim=1).reshape(3*self.embed_dim).contiguous()
+            else:
+                input_bias = self.in_proj_bias
+        else:
+            input_bias=None        
+        if key_padding_mask is not None:
+            assert (attn_mask is None), "ERROR attn_mask and key_padding_mask should not be both defined!"
+            mask = key_padding_mask
+        elif attn_mask is not None:
+            assert self.mask_additive == False, "additive mask not supported for time mask"
+            mask = attn_mask
+        else:
+            mask = None
+
+        if self.include_norm_add:
+            if self.impl == 'fast':
+                outputs = self.attn_func(attn_mask is not None, is_training, self.num_heads, query,
+                                         self.lyr_nrm_gamma_weights, self.lyr_nrm_beta_weights,
+                                         input_weights, self.out_proj_weight, mask, self.dropout)
+            else:
+                lyr_nrm_results = self.lyr_nrm(query)
+                outputs = self.attn_func(attn_mask is not None, is_training, self.num_heads, self.scaling, lyr_nrm_results,
+                                         input_weights, self.out_proj_weight,
+                                         input_bias, self.out_proj_bias,
+                                         mask, self.dropout)
+                if is_training:
+                    outputs = jit_dropout_add(outputs, query, self.dropout, is_training)
+                else:
+                    outputs = outputs + query
+        else:
+            if self.impl == 'fast':
+                outputs = self.attn_func(attn_mask is not None, is_training, self.num_heads, query,
+                                         input_weights, self.out_proj_weight, input_bias, self.out_proj_bias, mask, self.mask_additive, self.dropout)
+            else:
+                outputs = self.attn_func(attn_mask is not None, is_training, self.num_heads, self.scaling, query,
+                                         input_weights, self.out_proj_weight,
+                                         input_bias, self.out_proj_bias,
+                                         mask, self.mask_additive, self.dropout)
+
+        return outputs,None
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/self_multihead_attn_func.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/self_multihead_attn_func.py
new file mode 100644
index 0000000000..522dd446af
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/multihead_attn/self_multihead_attn_func.py
@@ -0,0 +1,235 @@
+import torch
+import torch.nn.functional as F
+
+class SelfAttnFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, use_time_mask, is_training, heads, scale, inputs,
+                input_weights, output_weights,
+                input_biases, output_biases,
+                mask, is_additive_mask, dropout_prob):
+        use_biases_t   = torch.tensor([input_biases is not None])
+        heads_t        = torch.tensor([heads])
+        scale_t        = torch.tensor([scale])
+        dropout_prob_t = torch.tensor([dropout_prob])
+        null_tensor    = torch.tensor([])
+        head_dim       = inputs.size(2) // heads
+
+        # Input Linear GEMM
+        # input1: (activations) [seql_q, seqs, embed_dim(1024)]
+        # input2: (weights)     [embed_dim*3 (3072), embed_dim (1024)] (transpose [0,1])
+        # output:               [seql_q, seqs, embed_dim*3]
+        # GEMM: ( (seql_q*seqs) x embed_dim ) x ( embed_dim x embed_dim*3 ) = (seql_q*seqs x embed_dim*3)
+        if use_biases_t[0]:
+            input_lin_results = torch.addmm(input_biases,
+                                            inputs.view(inputs.size(0) * inputs.size(1), inputs.size(2)),
+                                            input_weights.transpose(0,1),
+                                            beta=1., alpha=1.)
+        else:
+            input_lin_results = torch.mm(inputs.view(inputs.size(0) * inputs.size(1), inputs.size(2)), input_weights.transpose(0,1))
+        input_lin_results = input_lin_results.view(inputs.size(0), inputs.size(1), input_weights.size(0))
+
+        # Slice out q,k,v from one big Input Linear outuput (should only impact meta data, no copies!)
+        # Sequences and heads are combined to make the batch of the Batched GEMM
+        # input_lin_results: [seql_q, seqs, heads(16), 3, head_dim(64)]
+        # input_lin_results: [seql_q, batches=seqs*heads, 3, head_dim]
+        input_lin_results = input_lin_results.view(inputs.size(0), inputs.size(1)*heads, 3, head_dim)
+        queries = input_lin_results[:,:,0,:]
+        keys    = input_lin_results[:,:,1,:]
+        values  = input_lin_results[:,:,2,:]
+
+        # Matmul1 Batched GEMMs
+        # The output tensor is specified prior to the Batch GEMM because baddbmm requires its specification
+        # baddbmm is used to apply the scale parameter via the Batched GEMM's alpha parameter instead of
+        # a separate elementwise operation.
+        # Input1: (Queries) [seql_q, seqs*heads, head_dim] tranpose(0,1)
+        # Input2: (Keys)    [seql_k, seqs*heads, head_dim] transpose(0,1)
+        # output:           [seqs*heads, seql_q, seql_k]
+        # GEMM: Per batch: ( seql_q x head_dim ) x ( head_dim x seql_k ) = ( seql_q x seql_k )
+        matmul1_results = torch.empty((queries.size(1),queries.size(0),keys.size(0)), dtype=queries.dtype, device=torch.device('cuda'))
+        matmul1_results = torch.baddbmm(matmul1_results, queries.transpose(0,1), keys.transpose(0,1).transpose(1,2), out=matmul1_results, beta=0.0, alpha=scale_t[0])
+
+        if mask is not None:
+            # Self Attention Time Mask
+            if use_time_mask:
+                assert (len(mask.size()) == 2), "Timing mask is not 2D!"
+                assert (mask.size(0) == mask.size(1)), "Sequence length should match!"
+                mask = mask.to(torch.bool)
+                matmul1_results = matmul1_results.masked_fill_(mask, float('-inf'))
+            # Key Padding Mask
+            else:
+                batches,seql_q,seql_k = matmul1_results.size()
+                seqs = int(batches / heads)
+                matmul1_results = matmul1_results.view(seqs, heads, seql_q, seql_k)
+                if is_additive_mask:
+                    matmul1_results = matmul1_results + mask.unsqueeze(1).unsqueeze(2)
+                else:
+                    mask = mask.to(torch.bool)
+                    matmul1_results = matmul1_results.masked_fill_(mask.unsqueeze(1).unsqueeze(2), float('-inf'))
+                matmul1_results = matmul1_results.view(seqs*heads, seql_q, seql_k)
+
+        softmax_results = F.softmax(matmul1_results, dim=-1)
+
+        # Dropout - is not executed for inference
+        if is_training:
+            dropout_results,dropout_mask = torch._fused_dropout(softmax_results, p=(1.-dropout_prob_t[0]))
+        else:
+            dropout_results = softmax_results
+            dropout_mask    = null_tensor
+
+        # Matmul2 Batched GEMMs
+        # The output tensor specification is needed here to specify the non-standard output.
+        # Given that pytorch cannot currently perform autograd with an output tensor specified,
+        # this requires a backward pass specified.
+        # Input1: from_softmax [seqs*heads, seql_q, seql_k]
+        # Input2: (values)     [seql_v, seqs*heads, head_dim] transpose(0,1)
+        # Output:              [seql_q, seqs*heads, head_dim] transpose(0,1)
+        # GEMM: Per batch: ( seql_q x seql_k ) x ( seql_k x head_dim ) = (seql_q x head_dim)
+        matmul2_results = torch.empty((dropout_results.size(1), dropout_results.size(0), values.size(2)), dtype=dropout_results.dtype, device=torch.device('cuda')).transpose(1,0)
+        matmul2_results = torch.bmm(dropout_results, values.transpose(0,1), out=matmul2_results)
+        matmul2_results = matmul2_results.transpose(0, 1).contiguous().view(inputs.size(0), inputs.size(1), inputs.size(2))
+
+        # Output Linear GEMM
+        # Input1: (activations) [seql_q, seqs, embed_dim=heads*head_dim]
+        # Input2: (weights)     [ embed_dim, embed_dim ] transpose(0,1)
+        # Output:               [ seql_q, seqs, embed_dim ]
+        # GEMM: ( seql_q*seqs x embed_dim ) x ( embed_dim x embed_dim ) = ( seql_q*seqs x embed_dim )
+        if use_biases_t[0]:
+            outputs = torch.addmm(output_biases,
+                                  matmul2_results.view(inputs.size(0) * inputs.size(1), inputs.size(2)),
+                                  output_weights.transpose(0,1),
+                                  beta=1., alpha=1.)
+        else:
+            outputs = torch.mm(matmul2_results.view(inputs.size(0) * inputs.size(1), inputs.size(2)), output_weights.transpose(0,1))
+        outputs = outputs.view(inputs.size(0), inputs.size(1), output_weights.size(0))
+
+        ctx.save_for_backward(use_biases_t,                             \
+                              heads_t,                                  \
+                              scale_t,                                  \
+                              matmul2_results,                          \
+                              dropout_results,                          \
+                              softmax_results,                          \
+                              input_lin_results,                        \
+                              inputs,                                   \
+                              input_weights,                            \
+                              output_weights,                           \
+                              dropout_mask,                             \
+                              dropout_prob_t)
+
+        return outputs.detach()
+
+    @staticmethod
+    def backward(ctx, output_grads):
+        use_biases_t,                                                   \
+        heads_t,                                                        \
+        scale_t,                                                        \
+        matmul2_results,                                                \
+        dropout_results,                                                \
+        softmax_results,                                                \
+        input_lin_results,                                              \
+        inputs,                                                         \
+        input_weights,                                                  \
+        output_weights,                                                 \
+        dropout_mask,                                                   \
+        dropout_prob_t          = ctx.saved_tensors
+
+        head_dim                = inputs.size(2) // heads_t[0]
+
+        # Slice out q,k,v from one big Input Linear outuput (should only impact meta data, no copies!)
+        # Sequences and heads are combined to make the batch of the Batched GEMM
+        # input_lin_results: [seql_q, seqs, heads(16), 3, head_dim(64)]
+        # input_lin_results: [seql_q, batches=seqs*heads, 3, head_dim]
+        input_lin_results       = input_lin_results.view(inputs.size(0), inputs.size(1)*heads_t[0], 3, head_dim)
+        queries                 = input_lin_results[:,:,0,:]
+        keys                    = input_lin_results[:,:,1,:]
+        values                  = input_lin_results[:,:,2,:]
+
+        # Slice out q,k,v from one big set of gradients entering the input linear's bprop  (should only impact meta data, no copies!)
+        # The gradients are identical in size to the Input Linear outputs.
+        # The tensor is declared before hand to properly slice out query, key, and value grads.
+        input_lin_results_grads = torch.empty_like(input_lin_results)
+        queries_grads           = input_lin_results_grads[:,:,0,:]
+        keys_grads              = input_lin_results_grads[:,:,1,:]
+        values_grads            = input_lin_results_grads[:,:,2,:]
+
+        # Output Linear GEMM - DGRAD
+        # Input1: (data grads)  [seql_q, seqs, embed_dim=heads*head_dim]
+        # Input2: (weights)     [ embed_dim, embed_dim ]
+        # Output:               [ seql_q, seqs, embed_dim ]
+        # GEMM: ( seql_q*seqs x embed_dim ) x ( embed_dim x embed_dim ) = ( seql_q*seqs x embed_dim )
+        output_lin_grads = torch.mm(output_grads.view(output_grads.size(0) * output_grads.size(1), output_grads.size(2)), output_weights)
+        output_lin_grads = output_lin_grads.view(output_grads.size(0), output_grads.size(1), output_weights.size(1))
+        # Output Linear GEMM - WGRAD
+        # Input1: (data grads)  [seql_q*seqs, embed_dim=heads*head_dim] transpose(0,1)
+        # Input2: (activations) [seql_q*seqs, embed_dim ]
+        # Output:               [ seql_q, seqs, embed_dim ]
+        # GEMM: ( embed_dim x seql_q*seqs ) x ( seql_q*seqs x embed_dim ) = ( embed_dim x embed_dim )
+        output_weight_grads = torch.mm(output_grads.view(output_grads.size(0) * output_grads.size(1), output_grads.size(2)).transpose(0,1),
+                                       matmul2_results.view(matmul2_results.size(0) * matmul2_results.size(1), matmul2_results.size(2)))
+        output_lin_grads = output_lin_grads.view(inputs.size(0), inputs.size(1)*heads_t[0], head_dim).transpose(0,1)
+
+        if use_biases_t[0]:
+            output_bias_grads = torch.sum(output_grads.view(output_grads.size(0) * output_grads.size(1), output_grads.size(2)), 0)
+        else:
+            output_bias_grads = None
+
+        # Matmul2 - DGRAD1
+        # Input1: (data grads)  [seql_q, seqs*heads, head_dim] transpose(0,1)
+        # Input2: (activations) [seql_k, seqs*heads, head_dim] transpose(0,1).transpose(1,2)
+        # Output:               [seqs*heads, seql_q, seql_k]
+        # GEMM: Per batch: ( seql_q x head_dim ) x ( head_dim x seql_k ) = ( seql_q x seql_k )
+        matmul2_dgrad1 = torch.bmm(output_lin_grads, values.transpose(0,1).transpose(1,2))
+        # Matmul2 - DGRAD2
+        # Input1: (data grads)  [seql_q, seqs*heads, head_dim] transpose(0,1)
+        # Input2: (activations) [seql_k, seqs*heads, head_dim] transpose(0,1).transpose(1,2)
+        # Output:               [seqs*heads, seql_q, seql_k]
+        # GEMM: Per batch: ( seql_q x head_dim ) x ( head_dim x seql_k ) = ( seql_q x seql_k )
+        values_grads   = torch.bmm(dropout_results.transpose(1,2), output_lin_grads, out=values_grads.transpose(0,1))
+
+        # Mask and Scaling for Dropout (not a publically documented op)
+        dropout_grads = torch._masked_scale(matmul2_dgrad1, dropout_mask, 1.0/(1.0-dropout_prob_t[0]))
+
+        # Softmax Grad (not a publically documented op)
+        softmax_grads = torch._softmax_backward_data(dropout_grads, softmax_results, -1, softmax_results)
+
+        # Matmul1 - DGRAD1
+        # Input1: (data grads)  [seqs*heads, seql_q, seql_k] 
+        # Input2: (activations) [seql_k, seqs*heads, head_dim] transpose(0,1)
+        # Output:               [seqs*heads, seql_q, head_dim] transpose(0,1)
+        # GEMM: Per batch: ( seql_q x seql_k ) x ( seql_k x head_dim ) = ( seql_q x head_dim )
+        queries_grads = torch.baddbmm(queries_grads.transpose(0,1), softmax_grads, keys.transpose(0,1),
+                                      out=queries_grads.transpose(0,1), beta=0.0, alpha=scale_t[0])
+        # Matmul1 - DGRAD2
+        # Input1: (data grads)  [seqs*heads, seql_q, seql_k] transpose(1,2)
+        # Input2: (activations) [seql_q, seqs*heads, head_dim] transpose(0,1)
+        # Output:               [seqs*heads, seql_k, head_dim] transpose(0,1)
+        # GEMM: Per batch: ( seql_k x seql_q ) x ( seql_q x head_dim ) = ( seql_k x head_dim )
+        keys_grads    = torch.baddbmm(keys_grads.transpose(0,1), softmax_grads.transpose(1,2), queries.transpose(0,1),
+                                      out=keys_grads.transpose(0,1), beta=0.0, alpha=scale_t[0])
+
+        # Input Linear GEMM - DGRAD
+        # input1: (data grads) [seql_q, seqs, 3*embed_dim(3072)]
+        # input2: (weights)    [embed_dim*3 (3072), embed_dim (1024)] 
+        # output:              [seql_q, seqs, embed_dim]
+        # GEMM: ( (seql_q*seqs) x 3*embed_dim ) x ( 3*embed_dim x embed_dim ) = (seql_q*seqs x embed_dim)
+        input_lin_results_grads = input_lin_results_grads.view(inputs.size(0)*inputs.size(1), heads_t[0]*3*head_dim)
+        input_grads = torch.mm(input_lin_results_grads, input_weights)
+        input_grads = input_grads.view(inputs.size(0), inputs.size(1), inputs.size(2))
+        # Input Linear GEMM - WGRAD
+        # input1: (data grads)  [seql_q*seqs, 3*embed_dim(3072)]
+        # input2: (activations) [seql_q*seqs, embed_dim(1024)] 
+        # output:               [3*embed_dim, embed_dim]
+        # GEMM: ( 3*embed_dim x seql_q*seqs ) x ( seql_q*seqs x embed_dim ) = (3*embed_dim x embed_dim)
+        input_weight_grads = torch.mm(input_lin_results_grads.transpose(0,1), inputs.view(inputs.size(0)*inputs.size(1), inputs.size(2)))
+
+        if use_biases_t[0]:
+            input_bias_grads = torch.sum(input_lin_results_grads, 0)
+        else:
+            input_bias_grads = None
+
+        return None, None, None, None,                   \
+               input_grads,                              \
+               input_weight_grads, output_weight_grads,  \
+               input_bias_grads, output_bias_grads,      \
+               None, None
+
+self_attn_func = SelfAttnFunc.apply
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/__init__.py
new file mode 100644
index 0000000000..1933b2fdb6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/__init__.py
@@ -0,0 +1,3 @@
+from .fp16_optimizer import FP16_Optimizer
+from .fused_adam import FusedAdam
+from .fused_lamb import FusedLAMB
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam.py
new file mode 100644
index 0000000000..d9aa9bce7b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam.py
@@ -0,0 +1,636 @@
+import math
+import torch
+import importlib
+import amp_C
+from apex.multi_tensor_apply import multi_tensor_applier
+
+import torch.distributed.distributed_c10d as c10d
+
+class DistributedFusedAdam(torch.optim.Optimizer):
+
+    """Implements Adam algorithm. Currently GPU-only.  Requires Apex to be installed via
+    ``python setup.py install --cuda_ext --cpp_ext``.
+    
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+    
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square. (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        eps_inside_sqrt (boolean, optional): in the 'update parameters' step,
+            adds eps to the bias-corrected second moment estimate before
+            evaluating square root instead of adding it to the square root of
+            second moment estimate as in the original paper. (default: False)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False) NOT SUPPORTED in FusedAdam!
+        overlap_reductions(boolean, optional): whether to overlap reductions
+            with bprop (default: True)
+        step_supports_amp_scaling(boolean, optional): whether to use customized
+            gradient unscaling logic (default: True)
+        num_process_groups (integer, optional): number of process groups in
+            the app (default: 1)
+        current_process_group (object, optional): the process group to work on
+            (default: None)
+        process_group_id (integer, optional): process group id (default: 0)
+        process_group_size (integer, optional): size of process group
+            (default: 0)
+        clip_grad_norm (boolean, optional): whether to handle gradient clipping
+            (default: True)
+        model_parallel (boolean, optional): whether model parallelism is used
+            (default: False)
+
+
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params,
+                 lr=1e-3, bias_correction=True, betas=(0.9, 0.999),
+                 eps=1e-8, eps_inside_sqrt=False,
+                 weight_decay=0., max_grad_norm=0.,
+                 amsgrad=False, flat_mt=False,
+                 overlap_reductions=True,
+                 compute_L2_grad_norm=False,
+                 dwu_group_size=0, dwu_num_blocks=4, dwu_num_chunks=4,
+                 dwu_num_rs_pg=1, dwu_num_ar_pg=4, dwu_num_ag_pg=0,
+                 predivide=True, e5m2_allgather=False,
+                 do_not_flatten_model=False,
+                 step_supports_amp_scaling=True,
+                 num_process_groups=1,
+                 current_process_group=None,
+                 process_group_id=0,
+                 process_group_size=0,
+                 clip_grad_norm=True,
+                 model_parallel=False):
+        global fused_adam_cuda, distributed_adam_cuda
+        fused_adam_cuda = importlib.import_module("fused_adam_cuda")
+        distributed_adam_cuda = importlib.import_module("distributed_adam_cuda")
+        self.multi_tensor_l2norm = amp_C.multi_tensor_l2norm
+
+        if amsgrad:
+            raise RuntimeError('DistributedFusedAdam does not support the AMSGrad variant.')
+
+        defaults = dict(lr=lr, bias_correction=bias_correction,
+                        betas=betas, eps=eps, weight_decay=weight_decay,
+                        max_grad_norm=max_grad_norm)
+        super(DistributedFusedAdam, self).__init__(params, defaults)
+
+        # Misc
+        self.eps_mode = 0 if eps_inside_sqrt else 1
+        self._overflow_buf = torch.cuda.IntTensor([0])
+        self._has_overflow = False
+        self._step_supports_amp_scaling = step_supports_amp_scaling
+        self._last_step = False
+        self._overlap_reductions = overlap_reductions
+        self._global_scale = None
+        self._num_blocks = dwu_num_blocks
+        self._num_chunks = dwu_num_chunks
+        self._predivide = predivide
+        self._e5m2_allgather = e5m2_allgather
+        self._do_not_flatten_model = do_not_flatten_model
+        self._compute_L2_grad_norm = compute_L2_grad_norm
+        self._L2_grad_norm = None
+        self._flat_mt = flat_mt
+        self._init_done = False
+        self._resume_from_checkpoint = False
+        self._step = 0
+
+        # Process group related
+        self._clip_grad_norm = clip_grad_norm
+        self._model_parallel = model_parallel
+        self._num_process_groups = num_process_groups
+        self._current_process_group = current_process_group if current_process_group is not None else c10d._get_default_group()
+        self._available_ranks = list(c10d._pg_group_ranks[self._current_process_group].keys())
+        self._process_group_id = process_group_id
+        self._process_group_size = torch.cuda.device_count() if process_group_size <= 0 else process_group_size
+        self._world_size = self._process_group_size # world: the current process group
+        self._group_size = torch.cuda.device_count() if dwu_group_size <= 0 else dwu_group_size
+        self._num_groups = self._world_size // self._group_size
+        self._global_rank = torch.distributed.get_rank()
+        self._world_rank = self._global_rank // self._num_process_groups
+        self._group_rank = self._world_rank % self._group_size
+        #print("world_size:", self._world_size, ", group_size:", self._group_size, ", num_groups:", self._num_groups, ", global_rank:", self._global_rank, ", world_rank:", self._world_rank, ", group_rank:", self._group_rank)
+        self._num_rs_pg = dwu_num_rs_pg
+        self._num_ar_pg = dwu_num_ar_pg
+        self._num_ag_pg = dwu_num_ag_pg
+
+        # Master weight, moment, gradient buffers
+        self._fp32_p, self._fp32_m, self._fp32_v, self._fp16_p, self._fp16_g = None, None, None, None, None
+
+    def _first_step_init(self):
+        p_offset = 0
+        p_i = 0
+        self._model_params = []
+        self._grads_info = []
+        self._grad_accs = []
+        self._group_properties = []
+        for group in self.param_groups:
+            self._param_group = group
+            prev = None
+            beta1, beta2 = group['betas']
+            bias_correction = 1 if group['bias_correction'] else 0
+            eps = group['eps']
+            weight_decay = group['weight_decay']
+            for p in group['params']:
+                # broadcast from rank 0 of current process group
+                torch.distributed.broadcast(p, src=self._available_ranks[0], group=self._current_process_group)
+                if not p.requires_grad:
+                    continue
+                self._model_params.append(p)
+                # Multiple param groups support: 
+                # store one hyperparam item per parameter tensor
+                self._group_properties.append((
+                    beta1,
+                    beta2,
+                    bias_correction,
+                    eps,
+                    weight_decay
+                    ))
+                p_grads_size = p.numel()
+                def wrapper(param, param_i, param_grads_size, param_offset):
+                    param_tmp = param.expand_as(param)
+                    grad_acc = param_tmp.grad_fn.next_functions[0][0]
+                    def allreduce_hook(*unused):
+                        self._do_overlapped_reduction(param_i, param_grads_size, param_offset, param)
+                    grad_acc.register_hook(allreduce_hook)
+                    self._grad_accs.append(grad_acc)
+                self._grads_info.append({"param_grads_size":p_grads_size, "param_offset":p_offset})
+                wrapper(p, p_i, p_grads_size, p_offset)
+                p_offset += p_grads_size
+                # Only enforce 128b alignment (64 * fp16) for non-consecutive parameters
+                # RNN is one example of consecutive parameters:
+                # (weight_ih, weight_hh, bias_ih, bias_hh)
+                if prev is not None and (prev.data_ptr() + prev.numel() * prev.element_size() != p.data_ptr()):
+                    p_offset = ((p_offset + 63) // 64) * 64
+                prev = p
+                p_i += 1
+        self._grads_generated = [False]*len(self._grads_info)
+        self._grads = []
+        if self._overlap_reductions:
+            self._current_block = self._num_blocks
+
+        self._net_total_param_size = p_offset
+        self._total_param_size = p_offset
+        dwu_min_page_size = 256 * self._num_blocks * self._num_chunks * self._group_size
+        self._total_param_size = ((self._total_param_size + dwu_min_page_size - 1) // dwu_min_page_size) * dwu_min_page_size
+        self._block_size = self._total_param_size // self._num_blocks
+        self._chunk_size = self._block_size // self._num_chunks
+        self._shard_size = self._chunk_size // self._group_size
+        #print("self._net_total_param_size=%d, self._total_param_size=%d, dwu_min_page_size=%d, self._block_size=%d, self._chunk_size=%d, self._shard_size=%d" % (self._net_total_param_size, self._total_param_size,dwu_min_page_size,self._block_size,self._chunk_size,self._shard_size))
+
+        self._low_param_i = [0]*self._num_blocks
+        for block_id in range(self._num_blocks-1,-1,-1):
+            p_i = len(self._grads_info)-1
+            while p_i > 0 and self._grads_info[p_i]["param_offset"] > block_id*self._block_size:
+                p_i -= 1
+            self._low_param_i[block_id] = p_i
+        #print(self._low_param_i)
+
+        self._flat_grads = torch.zeros([self._total_param_size], dtype=torch.float16, device='cuda')
+        self._new_params = torch.zeros([self._total_param_size], dtype=torch.uint8 if self._e5m2_allgather else torch.float16, device='cuda')
+        self._mega_shard_size = self._num_blocks * self._num_chunks * self._shard_size
+        # initialize master weights, moments buffers if not loaded from checkpoint
+        if self._fp32_p is None:
+            self._fp32_p = torch.zeros([self._mega_shard_size], dtype=torch.float32, device='cuda')
+            self._fp32_m = torch.zeros([self._mega_shard_size], dtype=torch.float32, device='cuda')
+            self._fp32_v = torch.zeros([self._mega_shard_size], dtype=torch.float32, device='cuda')
+        # FIXME: Rethink fp16 label since it's either uint8 or fp16
+        self._fp16_p = torch.zeros([self._mega_shard_size], dtype=torch.uint8 if self._e5m2_allgather else torch.float16, device='cuda')
+        self._fp16_g = torch.zeros([self._mega_shard_size], dtype=torch.float16, device='cuda')
+
+        self._individual_flat_grads = []
+        for p_i, (grads_info, p) in enumerate(zip(self._grads_info, self._model_params)):
+            self._individual_flat_grads.append(self._flat_grads[grads_info["param_offset"]:grads_info["param_offset"]+grads_info["param_grads_size"]].view_as(p))
+
+        def _flat_split(p):
+            def __blockify(p):
+                return [p[block_id*self._block_size:(block_id+1)*self._block_size] for block_id in range(self._num_blocks)]
+            def __chunkify(p):
+                return [p[chunk_id*self._chunk_size:(chunk_id+1)*self._chunk_size] for chunk_id in range(self._num_chunks)]
+            def __shardify(p):
+                return [p[shard_id*self._shard_size:(shard_id+1)*self._shard_size] for shard_id in range(self._group_size)]
+            list_of_blocks = __blockify(self._flat_grads)
+            list_of_list_of_chunks = [__chunkify(block) for block in list_of_blocks]
+            list_of_list_of_list_of_shards = [[__shardify(chunk) for chunk in chunks] for chunks in list_of_list_of_chunks]
+            return list_of_blocks, list_of_list_of_chunks, list_of_list_of_list_of_shards
+        self._flat_grads_blocks, self._flat_grads_chunks, self._flat_grads_shards = _flat_split(self._flat_grads)
+        def _full_packed_split(p):
+            def __shardify(p):
+                return [p[mega_shard*self._mega_shard_size:(mega_shard+1)*self._mega_shard_size] for mega_shard in range(self._group_size)]
+            def __blockify(p):
+                return [p[block_id*self._num_chunks*self._shard_size:(block_id+1)*self._num_chunks*self._shard_size] for block_id in range(self._num_blocks)]
+            def __chunkify(p):
+                return [p[chunk_id*self._shard_size:(chunk_id+1)*self._shard_size] for chunk_id in range(self._num_chunks)]
+            list_of_mega_shards = __shardify(p)
+            list_of_list_of_mega_blocks = [__blockify(mega_shard) for mega_shard in list_of_mega_shards]
+            list_of_list_of_list_of_mega_chunks = [[__chunkify(mega_block) for mega_block in mega_blocks] for mega_blocks in list_of_list_of_mega_blocks]
+            return list_of_mega_shards, list_of_list_of_mega_blocks, list_of_list_of_list_of_mega_chunks
+        self._new_params_mega_shards, self._new_params_mega_blocks, self._new_params_mega_chunks = _full_packed_split(self._new_params)
+        def _packed_split(p):
+            def __packed_blockify(p):
+                packed_block_size = self._num_chunks*self._shard_size
+                return [p[block_id*packed_block_size:(block_id+1)*packed_block_size] for block_id in range(self._num_blocks)]
+            def __packed_chunkify(p):
+                # in the packed format, each chunk contains one shard, so packed_chunk_size == self._shard_size
+                return [p[chunk_id*self._shard_size:(chunk_id+1)*self._shard_size] for chunk_id in range(self._num_chunks)]
+            list_of_blocks = __packed_blockify(p)
+            list_of_list_of_chunks = [__packed_chunkify(block) for block in list_of_blocks]
+            return list_of_blocks, list_of_list_of_chunks
+        self._fp32_p_blocks, self._fp32_p_chunks = _packed_split(self._fp32_p)
+        self._fp32_m_blocks, self._fp32_m_chunks = _packed_split(self._fp32_m)
+        self._fp32_v_blocks, self._fp32_v_chunks = _packed_split(self._fp32_v)
+        self._fp16_p_blocks, self._fp16_p_chunks = _packed_split(self._fp16_p)
+        self._fp16_g_blocks, self._fp16_g_chunks = _packed_split(self._fp16_g)
+
+        # This paragraph does two things:
+        # 1) Copy model parameters into master buffer
+        # 2) Create tensor lists for unpacking new parameter tensor after all-gather
+        self._packed_flat_to_model_params = []
+        self._contrib_tensor_list = []
+        self._contrib_group_properties = []
+        self._non_parallel_grads = []
+        for shard_id in range(self._group_size):
+            for block_id in range(self._num_blocks):
+                for chunk_id in range(self._num_chunks):
+                    flat_shard_start = (((block_id * self._num_chunks + chunk_id) * self._group_size) + shard_id) * self._shard_size
+                    flat_shard_end = flat_shard_start + self._shard_size
+                    for (p, grads_info, group_props) in zip(self._model_params, self._grads_info, self._group_properties):
+                        flat_grad_start = grads_info["param_offset"]
+                        flat_grad_end = flat_grad_start + grads_info["param_grads_size"]
+                        clipped_start = (lambda a,b: a if a > b else b)(flat_grad_start, flat_shard_start)
+                        clipped_end = (lambda a,b: a if a < b else b)(flat_grad_end, flat_shard_end)
+                        if clipped_start < clipped_end:
+                            grad_offset = clipped_start - flat_grad_start
+                            grad_length = clipped_end - clipped_start
+                            shard_offset = clipped_start - flat_shard_start
+                            model_param_fragment = p.view(-1)[grad_offset:grad_offset+grad_length]
+                            new_param_packed_fragment = self._new_params_mega_chunks[shard_id][block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                            self._packed_flat_to_model_params.append( (new_param_packed_fragment, model_param_fragment) )
+                            if shard_id == self._group_rank:
+                                # copy model parameters into master buffer
+                                master_param_fragment = self._fp32_p_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                opti_state_m_fragment = self._fp32_m_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                opti_state_v_fragment = self._fp32_v_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                opti_state_g_fragment = self._fp16_g_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                opti_state_p_fragment = self._fp16_p_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                #print("model_param_fragment.size()=%s, new_param_packed_fragment.size()=%s, master_param_fragment.size()=%s" % (str(model_param_fragment.size()), str(new_param_packed_fragment.size()), str(master_param_fragment.size())))
+                                if not self._resume_from_checkpoint:
+                                    master_param_fragment.copy_(model_param_fragment)
+                                self._contrib_group_properties.append(group_props)
+                                self._contrib_tensor_list.append((master_param_fragment, opti_state_m_fragment, opti_state_v_fragment, opti_state_g_fragment, opti_state_p_fragment)) # p, m, v, g, p_copy
+                                if self._model_parallel and hasattr(p, 'model_parallel') and not p.model_parallel:
+                                    self._non_parallel_grads.append(opti_state_g_fragment)
+
+        p, m, v, g, p_copy = list(zip(*self._contrib_tensor_list))
+        self._contrib_tensor_list = [p, m, v, g, p_copy]
+
+        math_type = self._fp32_p.dtype
+        beta1, beta2, bias_correction, epsilon, decay = list(zip(*self._contrib_group_properties))
+        self._contrib_beta1 = torch.tensor(beta1, dtype=math_type, device='cuda')
+        self._contrib_beta2 = torch.tensor(beta2, dtype=math_type, device='cuda')
+        self._contrib_bias_correction = torch.tensor(bias_correction, dtype=torch.int, device='cuda')
+        self._contrib_epsilon = torch.tensor(epsilon, dtype=math_type, device='cuda')
+        self._contrib_weight_decay = torch.tensor(decay, dtype=math_type, device='cuda')
+
+        p_in, p_out = zip(*self._packed_flat_to_model_params)
+        self._packed_flat_to_model_params = [p_in, p_out]
+
+        if self._num_groups > 1:
+            self._ar_pg = []
+            for i in range(self._num_process_groups):
+                # gather global ranks of all members of the current process group
+                ranks = [i+k*self._num_process_groups for k in range(self._process_group_size)]
+                for j in range(self._group_size):
+                    ar_idx = [j+k*self._group_size for k in range(self._num_groups)]
+                    ar_rank = [ranks[k] for k in ar_idx]
+                    #if self._global_rank in ar_rank:
+                    #    print("group for all reduce, ranks:", ar_rank)
+                    for _ in range(self._num_ar_pg):
+                        grp = torch.distributed.new_group(ranks=ar_rank)
+                        if self._global_rank in ar_rank:
+                            self._ar_pg.append(grp)
+            self._ar_st = [torch.cuda.Stream() for _ in range(self._num_ar_pg)]
+            for ar_pg in self._ar_pg:
+                torch.distributed.all_reduce(self._overflow_buf,group=ar_pg)
+
+        self._rs_pg, rs_ranks = [],[]
+        for i in range(self._num_process_groups):
+            ranks = [i+k*self._num_process_groups for k in range(self._process_group_size)]
+            for j in range(self._num_groups):
+                rs_idx = [j*self._group_size+k for k in range(self._group_size)]
+                rs_rank = [ranks[k] for k in rs_idx]
+                #if self._global_rank in rs_rank:
+                #    print("group for reduce scatter, ranks:", rs_rank)
+                for _ in range(self._num_rs_pg):
+                    grp = torch.distributed.new_group(ranks=rs_rank)
+                    if self._global_rank in rs_rank:
+                        self._rs_pg.append(grp)
+                if self._compute_L2_grad_norm:
+                    l2_grad_norm_pg = torch.distributed.new_group(ranks=rs_rank)
+                    if self._global_rank in rs_rank:
+                        self._l2_grad_norm_pg = l2_grad_norm_pg
+                        torch.distributed.all_reduce(self._overflow_buf,group=self._l2_grad_norm_pg)
+        self._rs_st = [torch.cuda.Stream() for _ in range(self._num_rs_pg)]
+        for rs_pg in self._rs_pg:
+            torch.distributed.all_reduce(self._overflow_buf,group=rs_pg)
+
+        if self._num_ag_pg == 0:
+            self._ag_pg = self._rs_pg
+            self._ag_st = self._rs_st
+            self._num_ag_pg = self._num_rs_pg
+        else:
+            self._ag_pg = []
+            for i in range(self._num_process_groups):
+                ranks = [i+k*self._num_process_groups for k in range(self._process_group_size)]
+                for j in range(self._num_groups):
+                    ag_rank = rs_ranks[j]
+                    #if self._global_rank in ag_rank:
+                    #    print("group for all gather, ranks:", ag_rank)
+                    for _ in range(self._num_ag_pg):
+                        grp = torch.distributed.new_group(ranks=ag_rank)
+                        if self._global_rank in ag_rank:
+                            self._ag_pg.append(grp)
+            self._ag_st = [torch.cuda.Stream() for _ in range(self._num_ag_pg)]
+            for ag_pg in self._ag_pg:
+                torch.distributed.all_reduce(self._overflow_buf,group=ag_pg)
+        self._l2_grad_norm_st = torch.cuda.Stream() if self._compute_L2_grad_norm else None
+        self._completion_st = torch.cuda.Stream()
+
+        self._reductions_works = [None]*self._num_blocks
+        self._allgather_works = [None]*self._num_blocks
+
+        import inspect
+        assert ('no_copy' in inspect.getfullargspec(torch.distributed.reduce_scatter).args), "This version of c10d does not support no_copy option"
+
+    def _init_everything(self):
+        if not self._init_done:
+            self._first_step_init()
+            self._init_done = True
+
+    def set_last_step(self, last_step):
+        self._last_step = last_step
+        
+    def _get_flush_block(self):
+        flush_block = []
+        if self._current_block > 0 and self._grads_generated[self._low_param_i[self._current_block-1]]:
+            num_grads = len(self._grads_generated)
+            contiguous_idx = num_grads
+            while contiguous_idx > 0 and self._grads_generated[contiguous_idx-1]:
+                contiguous_idx -= 1
+
+            if contiguous_idx < num_grads and self._grads_info[contiguous_idx]["param_offset"] <= (self._current_block-1)*self._block_size:
+                self._current_block -= 1
+                start = self._current_block * self._block_size
+                end = (self._current_block+1) * self._block_size
+                flush_block = [start, end]
+
+        return flush_block
+
+    def _pipeline_block_reductions(self, block_id):
+        self._flatten_grad_mt(1.0/self._world_size if self._predivide else 1.0)
+
+        # Reduction within each node
+        # Changes gradient format from [block * chunk * shard] to [shard * block * chunk]
+        # The output format is the same as the fp32 master parameters
+        works = [None]*self._num_chunks
+        for chunk_id in range(self._num_chunks):
+            glob_chunk_id = block_id * self._num_chunks + chunk_id
+            rs_stream = self._rs_st[glob_chunk_id%self._num_rs_pg]
+            rs_stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(rs_stream):
+                works[chunk_id] = torch.distributed.reduce_scatter(self._fp16_g_chunks[block_id][chunk_id],self._flat_grads_shards[block_id][chunk_id],group=self._rs_pg[glob_chunk_id%self._num_rs_pg],async_op=True,no_copy=True)
+
+        # Reduction across nodes for each rank
+        if self._num_groups > 1:
+            for chunk_id in range(self._num_chunks):
+                glob_chunk_id = block_id * self._num_chunks + chunk_id
+                ar_stream = self._ar_st[glob_chunk_id%self._num_ar_pg]
+                with torch.cuda.stream(ar_stream):
+                    works[chunk_id].wait()
+                    works[chunk_id] = torch.distributed.all_reduce(self._fp16_g_chunks[block_id][chunk_id],group=self._ar_pg[glob_chunk_id%self._num_ar_pg],async_op=True)
+        self._reductions_works[block_id] = works
+
+        # Optionally compute L2 grad norm
+        if self._compute_L2_grad_norm and block_id == 0:
+            with torch.cuda.stream(self._l2_grad_norm_st):
+                for block_id in range(self._num_blocks):
+                    for chunk_id in range(self._num_chunks):
+                        self._reductions_works[block_id][chunk_id].wait()
+                # Since the packed format is contiguous after reductions, only one norm is needed
+                l2_grad_norm_sq = torch.empty([1], device='cuda')
+                l2_grad_norm_sq = self._fp16_g.norm(dtype=torch.float32, p=2)**2
+                torch.distributed.all_reduce(l2_grad_norm_sq, group=self._l2_grad_norm_pg)
+                # for model_parallel_rank=0, keep all gradients
+                # for the rest, subtract non_parallel gradients
+                if self._model_parallel and self._process_group_id: # non zero model_parallel_rank
+                    non_parallel_grad_norm_sq = torch.zeros([1], device='cuda')
+                    if len(self._non_parallel_grads): # non parallel grads exit
+                        non_parallel_grad_norm_sq = multi_tensor_applier(self.multi_tensor_l2norm,
+                                                                         self._overflow_buf,
+                                                                         [self._non_parallel_grads], False)[0]**2
+                    torch.distributed.all_reduce(non_parallel_grad_norm_sq, group=self._l2_grad_norm_pg)
+                    l2_grad_norm_sq = l2_grad_norm_sq - non_parallel_grad_norm_sq
+                self._L2_grad_norm = l2_grad_norm_sq.sqrt().item()
+
+    def __launch_step_kernel(self):
+        # If self._clip_grad_norm is False, we assume gradient clipping already 
+        # happened outside the optimizer and self._global_scale has already 
+        # been set to the combined scale, i.e. it's no longer the current loss
+        # scale used by the loss scaler. 
+        # For model parallelism cases in which we need to get global gradient 
+        # norm via all-reduce outside the optimizer to do the clipping. 
+        combined_scale = self._global_scale
+        if self._clip_grad_norm and self._param_group['max_grad_norm'] > 0 and math.isfinite(self.L2_grad_norm):
+            combined_scale = self._param_group['max_grad_norm'] / (self.L2_grad_norm / self._global_scale + 1e-6)
+            combined_scale = self._global_scale / min(1, combined_scale)
+        
+        self._step += 1
+        multi_tensor_applier(distributed_adam_cuda.multi_tensor_fused_adam,
+                self._overflow_buf,
+                self._contrib_tensor_list, # p, m, v, g, p_copy
+                self._contrib_beta1,
+                self._contrib_beta2,
+                self._contrib_bias_correction,
+                self._contrib_epsilon,
+                self._contrib_weight_decay,
+                self._param_group['lr'],
+                combined_scale,
+                self._step,
+                self.eps_mode)
+
+    def _pipeline_step(self):
+        # Call step kernel once per step
+        # Call all-gather once per step
+        with torch.cuda.stream(self._completion_st):
+            for block_id in range(self._num_blocks):
+                for chunk_id in range(self._num_chunks):
+                    self._reductions_works[block_id][chunk_id].wait()
+            self.__launch_step_kernel()
+            torch.distributed.all_gather(self._new_params_mega_shards, self._fp16_p, group=self._ag_pg[0], no_copy=True)
+
+    def _flatten_grad_mt(self, scale):
+        if self._flat_mt and len(self._grads) > 0:
+            self._overflow_buf.zero_()
+            multi_tensor_applier(
+                    amp_C.multi_tensor_scale,
+                    self._overflow_buf,
+                    list(zip(*self._grads)),
+                    scale)
+            self._grads = []
+
+    def _do_overlapped_reduction(self, param_i, param_grads_size, param_offset, param):
+        # handle overlapped reductions
+        if self._flat_mt:
+            self._grads.append( (param.grad, self._individual_flat_grads[param_i]) )
+        else:
+            torch.div(param.grad, self._world_size if self._predivide else 1.0, out=self._individual_flat_grads[param_i])
+        self._grads_generated[param_i]=True
+        if not self._last_step:
+            if self._overlap_reductions:
+                flush_block = self._get_flush_block()
+                while flush_block:
+                    block_id = flush_block[0] // self._block_size
+                    self._pipeline_block_reductions(block_id)
+                    flush_block = self._get_flush_block()
+
+    def set_global_scale(self, global_scale):
+        """Set global scale.
+        """
+        self._global_scale = global_scale
+
+    @property
+    def global_scale(self):
+        return self._global_scale
+
+    @property
+    def has_overflow(self):
+        """Check if overflows were detected by any call to step(...) method.
+        Clears the overflow flag.
+        """
+        has_overflow = self._has_overflow
+        self._has_overflow = False
+        return has_overflow
+
+    @property
+    def peek_overflow(self):
+        """Check if overflows were detected by any call to step(...) method.
+        Does not clear overflow flag.
+        """
+        return self._has_overflow
+
+    def strided_check_finite(self, output_params, stride=1, start=-1, end=-1, clear=True):
+        """Strided check for overflow.
+        You can get status by calling has_overflow.
+        """
+        if start >= 0 and start < end:
+            out_p = output_params[start:end]
+        else:
+            out_p = output_params
+        fused_adam_cuda.strided_check_finite(self._overflow_buf,
+                out_p,
+                stride,
+                1 if clear else 0)
+        self._has_overflow = False if self._overflow_buf.item() == 0 else True
+        return self._has_overflow
+
+    @property
+    def L2_grad_norm(self):
+        if self._compute_L2_grad_norm:
+            torch.cuda.current_stream().wait_stream(self._l2_grad_norm_st)
+            return self._L2_grad_norm
+        else:
+            return None
+
+    def complete_reductions(self):
+        """Complete reductions if full pipeline is not selected or overlap is not allowed.
+        """
+        self._init_everything()
+        if self._last_step:
+            # zero out gradients that have not been completed yet
+            for param_i, grad_generated in enumerate(self._grads_generated):
+                if not grad_generated:
+                    grad_info = self._grads_info[param_i]
+                    param_offset = grad_info["param_offset"]
+                    param_size = grad_info["param_grads_size"]
+                    self._flat_grads[param_offset:param_offset+param_size].zero_()
+                    self._grads_generated[param_i] = True
+
+        if self._last_step or not self._overlap_reductions:
+            # nothing done so far, run full pipeline after reductions
+            for block_id in range(self._num_blocks-1,-1,-1):
+                self._pipeline_block_reductions(block_id)
+
+        if self._compute_L2_grad_norm:
+            torch.cuda.current_stream().wait_stream(self._l2_grad_norm_st)
+
+        self._current_block = self._num_blocks
+        self._grads_generated = [False]*len(self._grads_info)
+
+    def step(self, closure=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        self._pipeline_step()
+
+        with torch.cuda.stream(self._completion_st):
+            # Copy self._new_params to model params
+            multi_tensor_applier(
+                    fused_adam_cuda.maybe_cast_mt,
+                    self._overflow_buf,
+                    self._packed_flat_to_model_params)
+
+        torch.cuda.current_stream().wait_stream(self._completion_st)
+
+        self._reductions_works = [None]*self._num_blocks
+        self._allgather_works = [None]*self._num_blocks
+
+        return loss
+
+    def state_dict(self):
+        """
+        Returns a dict containing the current state of this :class:`DistributedFusedAdam` instance.
+        Example::
+            checkpoint = {}
+            checkpoint['model'] = model.state_dict()
+            checkpoint['optimizer'] = optimizer.state_dict()
+            torch.save(checkpoint, "saved.pth")
+        """
+        # save step, master weights and first/second moments
+        state_dict = {}
+        state_dict['step'] = self._step
+        state_dict['fp32_p'] = self._fp32_p
+        state_dict['fp32_m'] = self._fp32_m
+        state_dict['fp32_v'] = self._fp32_v
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+        """
+        Loads a state_dict created by an earlier call to state_dict().
+        If an DistributedFusedAdam instance was constructed from some ``init_optimizer``,
+        whose parameters in turn came from ``model``, it is expected that the user
+        will call ``model.load_state_dict()`` before
+        ``optimizer.load_state_dict()`` is called.
+        Example::
+            model = torch.nn.Linear(D_in, D_out).cuda().half()
+            optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
+            optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
+            ...
+            checkpoint = torch.load("saved.pth")
+            model.load_state_dict(checkpoint['model'])
+            optimizer.load_state_dict(checkpoint['optimizer'])
+        """
+        # restore step, master weights and first/second moments
+        self._step = state_dict['step']
+        self._fp32_p = state_dict['fp32_p'].to(device="cuda")
+        self._fp32_m = state_dict['fp32_m'].to(device="cuda")
+        self._fp32_v = state_dict['fp32_v'].to(device="cuda")
+        self._resume_from_checkpoint = True
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam_v2.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam_v2.py
new file mode 100644
index 0000000000..80a810835f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam_v2.py
@@ -0,0 +1,615 @@
+import math
+import torch
+import importlib
+import amp_C
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class DistributedFusedAdamV2(torch.optim.Optimizer):
+
+    """Implements Adam algorithm. Currently GPU-only.  Requires Apex to be installed via
+    ``python setup.py install --cuda_ext --cpp_ext``.
+
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square. (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False) NOT SUPPORTED in FusedAdam!
+        eps_inside_sqrt (boolean, optional): in the 'update parameters' step,
+            adds eps to the bias-corrected second moment estimate before
+            evaluating square root instead of adding it to the square root of
+            second moment estimate as in the original paper. (default: False)
+        use_mt (boolean, optional): use multi tensor apply for lower launch
+            latency. (default: False)
+        overlap_reductions(boolean, optional): whether to overlap reductions
+            with bprop (default: True)
+        num_prestats (integer, optional): number of fp64 stats that will be
+            reduced during first fp16 gradient reduction block. 
+
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params,
+                 lr=1e-3, bias_correction = True,
+                 betas=(0.9, 0.999), eps=1e-8, eps_inside_sqrt = False,
+                 weight_decay=0., max_grad_norm=0., amsgrad=False, use_mt=False,
+                 amp_scale_adjustment=1.0, overlap_reductions=True, full_pipeline=True,
+                 compute_L2_grad_norm=False, distributed_weight_update=0,
+                 dwu_group_size=0, dwu_num_blocks=4, dwu_num_rs_pg=1, dwu_num_ar_pg=4,
+                 dwu_num_ag_pg=0, revert_method=1, flat_mt=False,
+                 dwu_num_chunks=4, predivide=True, e5m2_allgather=False,
+                 do_not_flatten_model=False):
+        global fused_adam_cuda
+        fused_adam_cuda = importlib.import_module("fused_adam_cuda")
+
+        self._amp_scale_adjustment = amp_scale_adjustment
+
+        if use_mt:
+            raise RuntimeError('DistributedFusedAdam does not support use_mt.')
+        if amsgrad:
+            raise RuntimeError('DistributedFusedAdam does not support the AMSGrad variant.')
+
+        defaults = dict(lr=lr, bias_correction=bias_correction,
+                        betas=betas, eps=eps, weight_decay=weight_decay,
+                        max_grad_norm=max_grad_norm)
+        super(DistributedFusedAdamV2, self).__init__(params, defaults)
+        self.eps_mode = 0 if  eps_inside_sqrt else 1
+
+        self._overflow_buf = torch.cuda.IntTensor([0])
+        self._has_overflow = False
+
+        assert (len(self.param_groups) == 1), "More than one parameter group is not supported."
+
+        # Way to revert a step
+        # 3 -> undo kernel + double buffer (debug, print norm of difference)
+        # 2 -> double buffer fp32 parameters
+        # 1 -> undo kernel
+        self._revert_method = revert_method
+        if self._revert_method > 1:
+            print("revert_method -> double buffer fp32 parameters, will consume more memory")
+
+        self._last_step = False
+        self._overlap_reductions = overlap_reductions
+        self._global_scale = None
+        self._num_blocks = dwu_num_blocks
+        self._num_chunks = dwu_num_chunks
+        self._predivide = predivide
+        self._e5m2_allgather = e5m2_allgather
+        self._do_not_flatten_model = do_not_flatten_model
+        self._full_pipeline = full_pipeline
+        self._compute_L2_grad_norm = compute_L2_grad_norm
+        self._L2_grad_norm = None
+        self._group_size = torch.cuda.device_count() if dwu_group_size <= 0 else dwu_group_size
+        self._world_size = torch.distributed.get_world_size()
+        self._num_groups = self._world_size // self._group_size
+        self._rank_in_group = torch.distributed.get_rank() % self._group_size
+
+        p_offset = 0
+        p_i = 0
+        self._param_state = None
+        self._model_params = []
+        self._grads_info = []
+        self._grad_accs = []
+        for group in self.param_groups:
+            self._param_group = group
+            prev = None
+            for p in group['params']:
+                torch.distributed.broadcast(p,0)
+                if not p.requires_grad:
+                    continue
+                self._model_params.append(p)
+                state = self.state[p]
+                if len(state) == 0:
+                    state['step'] = 0
+                if self._param_state is None:
+                    self._param_state = state
+                p_grads_size = p.numel()
+                def wrapper(param, param_i, param_grads_size, param_offset):
+                    param_tmp = param.expand_as(param)
+                    grad_acc = param_tmp.grad_fn.next_functions[0][0]
+                    def allreduce_hook(*unused):
+                        self._do_overlapped_reduction(param_i, param_grads_size, param_offset, param)
+                    grad_acc.register_hook(allreduce_hook)
+                    self._grad_accs.append(grad_acc)
+                self._grads_info.append({"param_grads_size":p_grads_size, "param_offset":p_offset})
+                wrapper(p, p_i, p_grads_size, p_offset)
+                p_offset += p_grads_size
+                # Only enforce 128b alignment (64 * fp16) for non-consecutive parameters
+                # RNN is one example of consecutive parameters:
+                # (weight_ih, weight_hh, bias_ih, bias_hh)
+                if prev is not None and (prev.data_ptr() + prev.numel() * prev.element_size() != p.data_ptr()):
+                    p_offset = ((p_offset + 63) // 64) * 64
+                prev = p
+                p_i += 1
+        self._grads_generated = [False]*len(self._grads_info)
+        self._flat_mt = flat_mt
+        self._grads = []
+        if self._overlap_reductions:
+            self._current_block = self._num_blocks
+
+        self._net_total_param_size = p_offset
+        self._total_param_size = p_offset
+        dwu_min_page_size = 256 * self._num_blocks * self._num_chunks * self._group_size
+        self._total_param_size = ((self._total_param_size + dwu_min_page_size - 1) // dwu_min_page_size) * dwu_min_page_size
+        self._block_size = self._total_param_size // self._num_blocks
+        self._shard_size = self._block_size // self._group_size
+        self._chunk_size = self._shard_size // self._num_chunks
+        print("self._net_total_param_size=%d, self._total_param_size=%d, dwu_min_page_size=%d, self._block_size=%d, self._shard_size=%d, self._chunk_size=%d" % (self._net_total_param_size, self._total_param_size,dwu_min_page_size,self._block_size,self._shard_size,self._chunk_size))
+
+        self._low_param_i = [0]*self._num_blocks
+        for block_id in range(self._num_blocks-1,-1,-1):
+            p_i = len(self._grads_info)-1
+            while p_i > 0 and self._grads_info[p_i]["param_offset"] > block_id*self._block_size:
+                p_i -= 1
+            self._low_param_i[block_id] = p_i
+        print(self._low_param_i)
+
+        self._flat_grads = torch.zeros([self._total_param_size], dtype=torch.float16, device='cuda')
+        self._new_params = torch.zeros([self._total_param_size], dtype=torch.uint8 if self._e5m2_allgather else torch.float16, device='cuda')
+        self._mega_shard_size = self._num_blocks * self._num_chunks * self._chunk_size
+        self._fp32_p = torch.zeros([self._mega_shard_size], dtype=torch.float32, device='cuda')
+        self._fp32_m = torch.zeros([self._mega_shard_size], dtype=torch.float32, device='cuda')
+        self._fp32_v = torch.zeros([self._mega_shard_size], dtype=torch.float32, device='cuda')
+        # FIXME: Rethink fp16 label since it's either uint8 or fp16
+        self._fp16_p = torch.zeros([self._mega_shard_size], dtype=torch.uint8 if self._e5m2_allgather else torch.float16, device='cuda')
+        self._fp16_g = torch.zeros([self._mega_shard_size], dtype=torch.float16, device='cuda')
+
+        self._individual_flat_grads = []
+        for p_i, (grads_info, p) in enumerate(zip(self._grads_info, self._model_params)):
+            self._individual_flat_grads.append(self._flat_grads[grads_info["param_offset"]:grads_info["param_offset"]+grads_info["param_grads_size"]].view_as(p))
+
+        def _flat_split(p):
+            def __blockify(p):
+                return [p[block_id*self._block_size:(block_id+1)*self._block_size] for block_id in range(self._num_blocks)]
+            def __shardify(p):
+                return [p[shard_id*self._shard_size:(shard_id+1)*self._shard_size] for shard_id in range(self._group_size)]
+            def __chunkify(p):
+                return [p[chunk_id*self._chunk_size:(chunk_id+1)*self._chunk_size] for chunk_id in range(self._group_size)]
+            list_of_blocks = __blockify(self._flat_grads)
+            list_of_list_of_shards = [__shardify(block) for block in list_of_blocks]
+            list_of_list_of_list_of_chunks = [[__chunkify(shard) for shard in shards] for shards in list_of_list_of_shards]
+            return list_of_blocks, list_of_list_of_shards, list_of_list_of_list_of_chunks
+        self._flat_grads_blocks, self._flat_grads_shards, self._flat_grads_chunks = _flat_split(self._flat_grads)
+        def _full_packed_split(p):
+            def __shardify(p):
+                return [p[mega_shard*self._mega_shard_size:(mega_shard+1)*self._mega_shard_size] for mega_shard in range(self._group_size)]
+            def __blockify(p):
+                return [p[block_id*self._num_chunks*self._chunk_size:(block_id+1)*self._num_chunks*self._chunk_size] for block_id in range(self._num_blocks)]
+            def __chunkify(p):
+                return [p[chunk_id*self._chunk_size:(chunk_id+1)*self._chunk_size] for chunk_id in range(self._num_chunks)]
+            list_of_mega_shards = __shardify(p)
+            list_of_list_of_mega_blocks = [__blockify(mega_shard) for mega_shard in list_of_mega_shards]
+            list_of_list_of_list_of_mega_chunks = [[__chunkify(mega_block) for mega_block in mega_blocks] for mega_blocks in list_of_list_of_mega_blocks]
+            return list_of_mega_shards, list_of_list_of_mega_blocks, list_of_list_of_list_of_mega_chunks
+        self._new_params_mega_shards, self._new_params_mega_blocks, self._new_params_mega_chunks = _full_packed_split(self._new_params)
+        def _packed_split(p):
+            def __packed_blockify(p):
+                packed_block_size = self._num_chunks*self._chunk_size
+                return [p[block_id*packed_block_size:(block_id+1)*packed_block_size] for block_id in range(self._num_blocks)]
+            def __packed_chunkify(p):
+                return [p[chunk_id*self._chunk_size:(chunk_id+1)*self._chunk_size] for chunk_id in range(self._num_chunks)]
+            list_of_blocks = __packed_blockify(p)
+            list_of_list_of_chunks = [__packed_chunkify(block) for block in list_of_blocks]
+            return list_of_blocks, list_of_list_of_chunks
+        self._fp32_p_blocks, self._fp32_p_chunks = _packed_split(self._fp32_p)
+        self._fp32_m_blocks, self._fp32_m_chunks = _packed_split(self._fp32_m)
+        self._fp32_v_blocks, self._fp32_v_chunks = _packed_split(self._fp32_v)
+        self._fp16_p_blocks, self._fp16_p_chunks = _packed_split(self._fp16_p)
+        self._fp16_g_blocks, self._fp16_g_chunks = _packed_split(self._fp16_g)
+
+        # current arrangement
+        # 
+        # self._flat_grads
+        # self._flat_grads_blocks [x self._num_blocks, self._block_size]
+        # self._flat_grads_chunks [x self._num_chunks, self._chunk_size]
+        # self._flat_grads_shards [x self._group_size, self._shard_size]
+        #
+        # self._new_params
+        # self._new_params_mega_shards [x self._group_size, self._num_blocks*self._num_chunks*self._shard_size]
+        # self._new_params_mega_blocks [x self._num_blocks, self._num_chunks*self._shard_size]
+        # self._new_params_mega_chunks [x self._num_chunks, self._shard_size]
+        #
+        # self._fp32_p
+        # self._fp32_p_blocks [x self._num_blocks, self._num_chunks*self._shard_size]
+        # self._fp32_p_chunks [x self._num_chunks, self._shard_size]
+        # each chunk contains one shard
+        # same for self._fp32_m, self._fp32_v, self._fp16_p and self._fp16_g
+        #
+        # Usage:
+        # 
+        # for chunk_id in range(self._num_chunks):
+        #   works[chunk_id] = torch.distributed.reduce_scatter(self._flat_grads_chunks[block_id][chunk_id], self._fp16_g_chunks[block_id][chunk_id], ...)
+        #
+        # ----------------------------------------------------------------------------------------
+        #
+        # new arrangement
+        #
+        # NB! New equations for self._shard_size and self._chunk_size
+        #
+        # self._flat_grads
+        # self._flat_grads_blocks [x self._num_blocks, self._block_size]
+        # self._flat_grads_shards [x self._group_size, self._shard_size]
+        # self._flat_grads_chunks [x self._num_chunks, self._chunk_size]
+        #
+        # self._new_params
+        # self._new_params_mega_shards [x self._group_size, self._num_blocks*self._num_chunks*self._chunk_size]
+        # self._new_params_mega_blocks [x self._num_blocks, self._num_chunks*self._chunk_size]
+        # self._new_params_mega_chunks [x self._num_chunks, self._chunk_size]
+        #
+        # self._fp32_p
+        # self._fp32_p_blocks [x self._num_blocks, self._num_chunks*self._chunk_size]
+        # self._fp32_p_chunks [x self._num_chunks, self._chunk_size]
+        # same for self._fp32_m, self._fp32_v, self._fp16_p and self._fp16_g
+        #
+        # Usage:
+        #
+        # work = torch.distributed.reduce_scatter(self._flat_grads_blocks[block_id], self._fp16_g[block_id], ...)
+        # for chunk_id in range(self._num_chunks):
+        #   work.wait()
+        #   works[chunk_id] = torch.distributed.all_reduce(self._fp16_g_chunks[block_id][chunk_id], ...)
+        # or
+        # work.wait()
+        # works[0] = torch.distributed.all_reduce(self._fp16_g_blocks[block_id], ...)
+        #
+
+        # This paragraph does two things:
+        # 1) Copy model parameters into master buffer
+        # 2) Create tensor lists for unpacking new parameter tensor after all-gather
+        self._packed_flat_to_model_params = []
+        for shard_id in range(self._group_size):
+            for block_id in range(self._num_blocks):
+                flat_shard_start = (block_id * self._group_size + shard_id) * self._shard_size
+                flat_shard_end = flat_shard_start + self._shard_size
+                for p, grads_info in zip(self._model_params, self._grads_info):
+                    flat_grad_start = grads_info["param_offset"]
+                    flat_grad_end = flat_grad_start + grads_info["param_grads_size"]
+                    clipped_start = (lambda a,b: a if a > b else b)(flat_grad_start, flat_shard_start)
+                    clipped_end = (lambda a,b: a if a < b else b)(flat_grad_end, flat_shard_end)
+                    if clipped_start < clipped_end:
+                        grad_offset = clipped_start - flat_grad_start
+                        grad_length = clipped_end - clipped_start
+                        shard_offset = clipped_start - flat_shard_start
+                        model_param_fragment = p.view(-1)[grad_offset:grad_offset+grad_length]
+                        new_param_packed_fragment = self._new_params_mega_blocks[shard_id][block_id][shard_offset:shard_offset+grad_length]
+                        self._packed_flat_to_model_params.append( (new_param_packed_fragment, model_param_fragment) )
+                        if shard_id == self._rank_in_group:
+                            # copy model parameters into master buffer
+                            master_param_fragment = self._fp32_p_blocks[block_id][shard_offset:shard_offset+grad_length]
+                            print("model_param_fragment.size()=%s, new_param_packed_fragment.size()=%s, master_param_fragment.size()=%s" % (str(model_param_fragment.size()), str(new_param_packed_fragment.size()), str(master_param_fragment.size())))
+                            master_param_fragment.copy_(model_param_fragment)
+
+        p_in, p_out = zip(*self._packed_flat_to_model_params)
+        self._packed_flat_to_model_params = [p_in, p_out]
+
+        self._distributed_weight_update = distributed_weight_update # Is this still needed?
+        self._num_rs_pg = dwu_num_rs_pg
+        self._num_ar_pg = dwu_num_ar_pg
+        self._num_ag_pg = dwu_num_ag_pg
+        if self._num_groups > 1:
+            self._ar_pg = []
+            for dev_i in range(self._group_size):
+                ranks = [dev_i+j*self._group_size for j in range(self._num_groups)]
+                for i in range(self._num_ar_pg):
+                    grp = torch.distributed.new_group(ranks=ranks)
+                    if torch.distributed.get_rank() in ranks:
+                        self._ar_pg.append(grp)
+            self._ar_st = [torch.cuda.Stream() for _ in range(self._num_ar_pg)]
+            for ar_pg in self._ar_pg:
+                torch.distributed.all_reduce(self._overflow_buf,group=ar_pg)
+        rs_ranks = []
+        for group_i in range(self._num_groups):
+            rs_ranks.append([group_i*self._group_size+j for j in range(self._group_size)])
+        self._rs_pg = []
+        for group_i in range(self._num_groups):
+            ranks = rs_ranks[group_i]
+            for i in range(self._num_rs_pg):
+                grp = torch.distributed.new_group(ranks=ranks)
+                if torch.distributed.get_rank() in ranks:
+                    self._rs_pg.append(grp)
+            if self._compute_L2_grad_norm and torch.distributed.get_rank() in ranks:
+                self._l2_grad_norm_pg = torch.distributed.new_group(ranks=ranks)
+                torch.distributed.all_reduce(self._overflow_buf,group=self._l2_grad_norm_pg)
+        self._rs_st = [torch.cuda.Stream() for _ in range(self._num_rs_pg)]
+        for rs_pg in self._rs_pg:
+            torch.distributed.all_reduce(self._overflow_buf,group=rs_pg)
+        if self._num_ag_pg == 0:
+            self._ag_pg = self._rs_pg
+            self._ag_st = self._rs_st
+            self._num_ag_pg = self._num_rs_pg
+        else:
+            self._ag_pg = []
+            for group_i in range(self._num_groups):
+                ranks = rs_ranks[group_i]
+                for i in range(self._num_ag_pg):
+                    grp = torch.distributed.new_group(ranks=ranks)
+                    if torch.distributed.get_rank() in ranks:
+                        self._ag_pg.append(grp)
+            self._ag_st = [torch.cuda.Stream() for _ in range(self._num_ag_pg)]
+            for ag_pg in self._ag_pg:
+                torch.distributed.all_reduce(self._overflow_buf,group=ag_pg)
+        self._l2_grad_norm_st = torch.cuda.Stream() if self._compute_L2_grad_norm else None
+        self._completion_st = torch.cuda.Stream()
+
+        self._reductions_works = [None]*self._num_blocks
+        self._allgather_works = [None]*self._num_blocks
+
+        import inspect
+        assert ('no_copy' in inspect.getfullargspec(torch.distributed.reduce_scatter).args), "This version of c10d does not support no_copy option"
+
+
+    def set_last_step(self, last_step):
+        self._last_step = last_step
+        
+    def _get_flush_block(self):
+        flush_block = []
+        if self._current_block > 0 and self._grads_generated[self._low_param_i[self._current_block-1]]:
+            num_grads = len(self._grads_generated)
+            contiguous_idx = num_grads
+            while contiguous_idx > 0 and self._grads_generated[contiguous_idx-1]:
+                contiguous_idx -= 1
+
+            if contiguous_idx < num_grads and self._grads_info[contiguous_idx]["param_offset"] <= (self._current_block-1)*self._block_size:
+                self._current_block -= 1
+                start = self._current_block * self._block_size
+                end = (self._current_block+1) * self._block_size
+                flush_block = [start, end]
+
+        return flush_block
+
+    def _pipeline_block_reductions(self, block_id):
+        self._flatten_grad_mt(1.0/self._world_size if self._predivide else 1.0)
+
+        # Reduction within each node
+        # Changes gradient format from [block * chunk * shard] to [shard * block * chunk]
+        # The output format is the same as the fp32 master parameters
+        works = [None]*self._num_chunks
+        rs_stream = self._rs_st[block_id%self._num_rs_pg]
+        rs_stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(rs_stream):
+            rs_work = torch.distributed.reduce_scatter(self._fp16_g_blocks[block_id],self._flat_grads_shards[block_id],group=self._rs_pg[block_id%self._num_rs_pg],async_op=True,no_copy=True)
+            for chunk_id in range(self._num_chunks):
+                works[chunk_id] = rs_work
+
+        # Reduction across nodes for each rank
+        if self._num_groups > 1:
+            for chunk_id in range(self._num_chunks):
+                glob_chunk_id = block_id * self._num_chunks + chunk_id
+                ar_stream = self._ar_st[glob_chunk_id%self._num_ar_pg]
+                with torch.cuda.stream(ar_stream):
+                    rs_work.wait()
+                    works[chunk_id] = torch.distributed.all_reduce(self._fp16_g_chunks[block_id][chunk_id],group=self._ar_pg[glob_chunk_id%self._num_ar_pg],async_op=True)
+        self._reductions_works[block_id] = works
+
+        # Optionally compute L2 grad norm
+        if self._compute_L2_grad_norm and block_id == 0:
+            with torch.cuda.stream(self._l2_grad_norm_st):
+                for block_id in range(self._num_blocks):
+                    for chunk_id in range(self._num_chunks):
+                        self._reductions_works[block_id][chunk_id].wait()
+                # Since the packed format is contiguous after reductions, only one norm is needed
+                l2_grad_norm_sq = torch.empty([1], device='cuda')
+                l2_grad_norm_sq = self._fp16_g.norm(dtype=torch.float32, p=2)**2
+                torch.distributed.all_reduce(l2_grad_norm_sq, group=self._l2_grad_norm_pg)
+                self._L2_grad_norm = l2_grad_norm_sq.sqrt().item()
+
+    def __launch_step_kernel(self, p, p_copy, m, v, g):
+        combined_scale = self._global_scale
+        if self._param_group['max_grad_norm'] > 0 and math.isfinite(self.L2_grad_norm):
+            combined_scale = self._param_group['max_grad_norm'] / (self.L2_grad_norm / self._global_scale + 1e-6)
+            combined_scale = self._global_scale / min(1, combined_scale)
+        bias_correction = 1 if self._param_group['bias_correction'] else 0
+        beta1, beta2 = self._param_group['betas']
+        fused_adam_cuda.reversible_adam(
+                p, p_copy, m, v, g,
+                self._param_group['lr'],
+                beta1,
+                beta2,
+                self._param_group['eps'],
+                combined_scale,
+                self._param_state['step']+1,
+                self.eps_mode,
+                bias_correction,
+                self._param_group['weight_decay'])
+
+    def _pipeline_block_step(self, block_id):
+        # Call step kernel once per block
+        ag_stream = self._ag_st[block_id%self._num_ag_pg]
+        with torch.cuda.stream(ag_stream):
+            for chunk_id in range(self._num_chunks):
+                self._reductions_works[block_id][chunk_id].wait()
+            self.__launch_step_kernel(
+                self._fp32_p_blocks[block_id],
+                self._fp16_p_blocks[block_id],
+                self._fp32_m_blocks[block_id],
+                self._fp32_v_blocks[block_id],
+                self._fp16_g_blocks[block_id])
+        # Call all-gather once per step.
+        # FIXME: Determine which is faster, one all-gather per block or a single all-gather at end
+        if block_id == 0:
+            for other_ag_stream in self._ag_st:
+                self._completion_st.wait_stream(other_ag_stream)
+            with torch.cuda.stream(self._completion_st):
+                torch.distributed.all_gather(self._new_params_mega_shards, self._fp16_p, group=self._ag_pg[0], no_copy=True)
+
+    def _pipeline_step(self):
+        # Call step kernel once per step
+        # Call all-gather once per step
+        with torch.cuda.stream(self._completion_st):
+            for block_id in range(self._num_blocks):
+                for chunk_id in range(self._num_chunks):
+                    self._reductions_works[block_id][chunk_id].wait()
+            self.__launch_step_kernel(
+                self._fp32_p,
+                self._fp16_p,
+                self._fp32_m,
+                self._fp32_v,
+                self._fp16_g)
+            torch.distributed.all_gather(self._new_params_mega_shards, self._fp16_p, group=self._ag_pg[0], no_copy=True)
+
+    def _flatten_grad_mt(self, scale):
+        if self._flat_mt and len(self._grads) > 0:
+            self._overflow_buf.zero_()
+            multi_tensor_applier(
+                    amp_C.multi_tensor_scale,
+                    self._overflow_buf,
+                    list(zip(*self._grads)),
+                    scale)
+            self._grads = []
+
+    def _do_overlapped_reduction(self, param_i, param_grads_size, param_offset, param):
+        # handle overlapped reductions
+        if self._flat_mt:
+            self._grads.append( (param.grad, self._individual_flat_grads[param_i]) )
+        else:
+            torch.div(param.grad, self._world_size if self._predivide else 1.0, out=self._individual_flat_grads[param_i])
+        self._grads_generated[param_i]=True
+        if not self._last_step:
+            if self._overlap_reductions:
+                flush_block = self._get_flush_block()
+                while flush_block:
+                    block_id = flush_block[0] // self._block_size
+                    self._pipeline_block_reductions(block_id)
+                    if self._full_pipeline:
+                        self._pipeline_block_step(block_id)
+                    flush_block = self._get_flush_block()
+
+    def set_global_scale(self, global_scale):
+        """Set global scale.
+        """
+        self._global_scale = global_scale
+
+    @property
+    def global_scale(self):
+        return self._global_scale
+
+    @property
+    def has_overflow(self):
+        """Check if overflows were detected by any call to step(...) method.
+        Clears the overflow flag.
+        """
+        has_overflow = self._has_overflow
+        self._has_overflow = False
+        return has_overflow
+
+    @property
+    def peek_overflow(self):
+        """Check if overflows were detected by any call to step(...) method.
+        Does not clear overflow flag.
+        """
+        return self._has_overflow
+
+    def strided_check_finite(self, output_params, stride=1, start=-1, end=-1, clear=True):
+        """Strided check for overflow.
+        You can get status by calling has_overflow.
+        """
+        if start >= 0 and start < end:
+            out_p = output_params[start:end]
+        else:
+            out_p = output_params
+        fused_adam_cuda.strided_check_finite(self._overflow_buf,
+                out_p,
+                stride,
+                1 if clear else 0)
+        self._has_overflow = False if self._overflow_buf.item() == 0 else True
+        return self._has_overflow
+
+    @property
+    def L2_grad_norm(self):
+        if self._compute_L2_grad_norm:
+            torch.cuda.current_stream().wait_stream(self._l2_grad_norm_st)
+            return self._L2_grad_norm
+        else:
+            return None
+
+    def complete_reductions(self):
+        """Complete reductions if full pipeline is not selected or overlap is not allowed.
+        """
+
+        if self._last_step:
+            # zero out gradients that have not been completed yet
+            for param_i, grad_generated in enumerate(self._grads_generated):
+                if not grad_generated:
+                    grad_info = self._grads_info[param_i]
+                    param_offset = grad_info["param_offset"]
+                    param_size = grad_info["param_grads_size"]
+                    self._flat_grads[param_offset:param_offset+param_size].zero_()
+                    self._grads_generated[param_i] = True
+
+        if self._last_step or not self._overlap_reductions:
+            # nothing done so far, run full pipeline after reductions
+            for block_id in range(self._num_blocks-1,-1,-1):
+                self._pipeline_block_reductions(block_id)
+
+        if self._compute_L2_grad_norm:
+            torch.cuda.current_stream().wait_stream(self._l2_grad_norm_st)
+
+        self._current_block = self._num_blocks
+        self._grads_generated = [False]*len(self._grads_info)
+
+    def revert_step(self):
+        """Revert effect of previously calling partial_step.
+        """
+        # Call undo kernel once per step
+        combined_scale = self._global_scale
+        if self._param_group['max_grad_norm'] > 0 and math.isfinite(self.L2_grad_norm):
+            combined_scale = self._param_group['max_grad_norm'] / (self.L2_grad_norm / self._global_scale + 1e-6)
+            combined_scale = self._global_scale / min(1, combined_scale)
+        bias_correction = 1 if self._param_group['bias_correction'] else 0
+        beta1, beta2 = self._param_group['betas']
+        fused_adam_cuda.maybe_adam_undo(
+                    torch.empty([0]),
+                    self._fp32_p,
+                    self._fp32_m,
+                    self._fp32_v,
+                    self._fp16_g,
+                    self._param_group['lr'],
+                    beta1,
+                    beta2,
+                    self._param_group['eps'],
+                    combined_scale,
+                    self._param_state['step']+1,
+                    self.eps_mode,
+                    bias_correction,
+                    self._param_group['weight_decay'])
+
+    def step(self, closure=None, skip_overflow_check=False):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        if self._last_step or not self._overlap_reductions or not self._full_pipeline:
+            self._pipeline_step()
+
+        with torch.cuda.stream(self._completion_st):
+            # Check for overflow
+            # Store state for loss scaler calculation
+            has_overflow = False if skip_overflow_check else self.strided_check_finite(self._new_params, stride=self._shard_size, start=0, end=self._net_total_param_size)
+            if has_overflow:
+                self.revert_step()
+            else:
+                # Copy self._new_params to model params
+                for p in self._model_params: self.state[p]['step'] += 1
+                multi_tensor_applier(
+                        fused_adam_cuda.maybe_cast_mt,
+                        self._overflow_buf,
+                        self._packed_flat_to_model_params)
+
+        torch.cuda.current_stream().wait_stream(self._completion_st)
+
+        self._reductions_works = [None]*self._num_blocks
+        self._allgather_works = [None]*self._num_blocks
+
+        return loss
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam_v3.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam_v3.py
new file mode 100644
index 0000000000..1b37603320
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_adam_v3.py
@@ -0,0 +1,325 @@
+import math
+import torch
+import importlib
+import amp_C
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class DistributedFusedAdamV3(torch.optim.Optimizer):
+
+    """Implements Adam algorithm. Currently GPU-only.  Requires Apex to be installed via
+    ``python setup.py install --cuda_ext --cpp_ext``.
+
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square. (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False) NOT SUPPORTED in FusedAdam!
+        eps_inside_sqrt (boolean, optional): in the 'update parameters' step,
+            adds eps to the bias-corrected second moment estimate before
+            evaluating square root instead of adding it to the square root of
+            second moment estimate as in the original paper. (default: False)
+        use_mt (boolean, optional): use multi tensor apply for lower launch
+            latency. (default: False)
+        overlap_reductions(boolean, optional): whether to overlap reductions
+            with bprop (default: True)
+        num_prestats (integer, optional): number of fp64 stats that will be
+            reduced during first fp16 gradient reduction block. 
+
+    .. _Adam\: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params,
+                 lr=1e-3, bias_correction = True,
+                 betas=(0.9, 0.999), eps=1e-8, eps_inside_sqrt = False,
+                 weight_decay=0., max_grad_norm=0., amsgrad=False, use_mt=False,
+                 amp_scale_adjustment=1.0, overlap_reductions=True, full_pipeline=True,
+                 compute_L2_grad_norm=False, distributed_weight_update=0,
+                 dwu_group_size=0, dwu_num_blocks=4, dwu_num_rs_pg=1, dwu_num_ar_pg=4,
+                 dwu_num_ag_pg=0, revert_method=1, flat_mt=False,
+                 dwu_num_chunks=4, predivide=True, e5m2_allgather=False,
+                 do_not_flatten_model=False):
+        global fused_adam_cuda
+        fused_adam_cuda = importlib.import_module("fused_adam_cuda")
+
+        self._amp_scale_adjustment = amp_scale_adjustment
+
+        if use_mt:
+            raise RuntimeError('DistributedFusedAdam does not support use_mt.')
+        if amsgrad:
+            raise RuntimeError('DistributedFusedAdam does not support the AMSGrad variant.')
+
+        defaults = dict(lr=lr, bias_correction=bias_correction,
+                        betas=betas, eps=eps, weight_decay=weight_decay,
+                        max_grad_norm=max_grad_norm)
+        super(DistributedFusedAdamV3, self).__init__(params, defaults)
+        self.eps_mode = 0 if  eps_inside_sqrt else 1
+
+        self._overflow_buf = torch.cuda.IntTensor([0])
+
+        assert (len(self.param_groups) == 1), "More than one parameter group is not supported."
+
+        # Way to revert a step
+        # 3 -> undo kernel + double buffer (debug, print norm of difference)
+        # 2 -> double buffer fp32 parameters
+        # 1 -> undo kernel
+        self._revert_method = revert_method
+        if self._revert_method > 1:
+            print("revert_method -> double buffer fp32 parameters, will consume more memory")
+
+        self._last_step = False
+        self._overlap_reductions = overlap_reductions
+        self._global_scale = None
+        self._num_blocks = dwu_num_blocks
+        self._predivide = predivide
+        self._e5m2_allgather = e5m2_allgather
+        self._do_not_flatten_model = do_not_flatten_model
+        self._full_pipeline = full_pipeline
+        self._L2_grad_norm = None
+        self._group_size = torch.cuda.device_count() if dwu_group_size <= 0 else dwu_group_size
+        self._world_size = torch.distributed.get_world_size()
+        self._num_groups = self._world_size // self._group_size
+        self._rank_in_group = torch.distributed.get_rank() % self._group_size
+
+        p_offset = 0
+        p_i = 0
+        self._param_state = None
+        self._model_params = []
+        self._grads_info = []
+        self._grad_accs = []
+        for group in self.param_groups:
+            self._param_group = group
+            prev = None
+            for p in group['params']:
+                torch.distributed.broadcast(p,0)
+                if not p.requires_grad:
+                    continue
+                self._model_params.append(p)
+                state = self.state[p]
+                if len(state) == 0:
+                    state['step'] = 0
+                if self._param_state is None:
+                    self._param_state = state
+                p_grads_size = p.numel()
+                def wrapper(param, param_i, param_grads_size, param_offset):
+                    param_tmp = param.expand_as(param)
+                    grad_acc = param_tmp.grad_fn.next_functions[0][0]
+                    def allreduce_hook(*unused):
+                        self._do_overlapped_reduction(param_i, param_grads_size, param_offset, param)
+                    grad_acc.register_hook(allreduce_hook)
+                    self._grad_accs.append(grad_acc)
+                self._grads_info.append({"param_grads_size":p_grads_size, "param_offset":p_offset})
+                wrapper(p, p_i, p_grads_size, p_offset)
+                p_offset += p_grads_size
+                # Only enforce 128b alignment (64 * fp16) for non-consecutive parameters
+                # RNN is one example of consecutive parameters:
+                # (weight_ih, weight_hh, bias_ih, bias_hh)
+                if prev is not None and (prev.data_ptr() + prev.numel() * prev.element_size() != p.data_ptr()):
+                    p_offset = ((p_offset + 63) // 64) * 64
+                prev = p
+                p_i += 1
+        self._grads_generated = [False]*len(self._grads_info)
+        self._flat_mt = flat_mt
+        self._grads = []
+        self._current_block = self._num_blocks
+
+        self._net_total_param_size = p_offset
+        self._total_param_size = p_offset
+        dwu_min_page_size = 256 * self._num_blocks * self._group_size
+        self._total_param_size = ((self._total_param_size + dwu_min_page_size - 1) // dwu_min_page_size) * dwu_min_page_size
+        self._block_size = self._total_param_size // self._num_blocks
+        self._shard_size = self._total_param_size // self._group_size
+        print("self._net_total_param_size=%d, self._total_param_size=%d, dwu_min_page_size=%d, self._block_size=%d, self._shard_size=%d" % (self._net_total_param_size, self._total_param_size,dwu_min_page_size,self._block_size,self._shard_size))
+
+        self._low_param_i = [0]*self._num_blocks
+        for block_id in range(self._num_blocks-1,-1,-1):
+            p_i = len(self._grads_info)-1
+            while p_i > 0 and self._grads_info[p_i]["param_offset"] > block_id*self._block_size:
+                p_i -= 1
+            self._low_param_i[block_id] = p_i
+        print(self._low_param_i)
+
+        self._flat_grads = torch.zeros([self._total_param_size], dtype=torch.float16, device='cuda')
+        self._flat_params = torch.zeros_like(self._flat_grads)
+
+        def _flat_split(flat):
+            def __flat_blockify(flat):
+                return [flat[block_id*self._block_size:(block_id+1)*self._block_size] for block_id in range(self._num_blocks)]
+            def __flat_shardify(flat):
+                return [flat[shard_id*self._shard_size:(shard_id+1)*self._shard_size] for shard_id in range(self._group_size)]
+            return __flat_blockify(flat), __flat_shardify(flat)
+        self._flat_grads_blocks, self._flat_grads_shards = _flat_split(self._flat_grads)
+        self._flat_params_blocks, self._flat_params_shards = _flat_split(self._flat_params)
+
+        # master params
+        self._fp32_p = torch.zeros([self._shard_size], dtype=torch.float32, device='cuda')
+        self._fp32_m = torch.zeros([self._shard_size], dtype=torch.float32, device='cuda')
+        self._fp32_v = torch.zeros([self._shard_size], dtype=torch.float32, device='cuda')
+
+        # copy model params to flat_params and set_ model params to flat_params.
+        self._individual_flat_grads = []
+        with torch.no_grad():
+            for p, grads_info in zip(self._model_params, self._grads_info):
+                start = grads_info["param_offset"]
+                end = start + grads_info["param_grads_size"]
+                flat_p = self._flat_params[start:end].view_as(p)
+                flat_p.copy_(p)
+                p.set_(flat_p)
+                flat_grad = self._flat_grads[start:end]
+                self._individual_flat_grads.append(flat_grad)
+        self._fp32_p.copy_(self._flat_params_shards[self._rank_in_group].float())
+
+        self._dwu_st = torch.cuda.Stream()
+        self._l2_grad_norm_st = torch.cuda.Stream()
+        for group_i in range(self._num_groups):
+            ranks = [group_i*self._group_size+local_rank for local_rank in range(self._group_size)]
+            pg = torch.distributed.new_group(ranks=ranks)
+            if torch.distributed.get_rank() in ranks:
+                self._ag_pg = pg
+                torch.distributed.all_reduce(self._overflow_buf, group=self._ag_pg)
+
+        import inspect
+        assert ('no_copy' in inspect.getfullargspec(torch.distributed.reduce_scatter).args), "This version of c10d does not support no_copy option"
+
+    @property
+    def has_overflow(self):
+        return True if not self.L2_grad_norm is None and not math.isfinite(self.L2_grad_norm) else False
+
+    def set_last_step(self, last_step):
+        self._last_step = last_step
+        
+    def _get_flush_block(self):
+        flush_block = []
+        if self._current_block > 0 and self._grads_generated[self._low_param_i[self._current_block-1]]:
+            num_grads = len(self._grads_generated)
+            contiguous_idx = num_grads
+            while contiguous_idx > 0 and self._grads_generated[contiguous_idx-1]:
+                contiguous_idx -= 1
+
+            if contiguous_idx < num_grads and self._grads_info[contiguous_idx]["param_offset"] <= (self._current_block-1)*self._block_size:
+                self._current_block -= 1
+                start = self._current_block * self._block_size
+                end = (self._current_block+1) * self._block_size
+                flush_block = [start, end]
+
+        return flush_block
+
+    def __launch_step_kernel(self, p, p_copy, m, v, g):
+        combined_scale = self._global_scale
+        if self._param_group['max_grad_norm'] > 0 and math.isfinite(self.L2_grad_norm):
+            combined_scale = self._param_group['max_grad_norm'] / (self.L2_grad_norm / self._global_scale + 1e-6)
+            combined_scale = self._global_scale / min(1, combined_scale)
+        bias_correction = 1 if self._param_group['bias_correction'] else 0
+        beta1, beta2 = self._param_group['betas']
+        fused_adam_cuda.reversible_adam(
+                p, p_copy, m, v, g,
+                self._param_group['lr'],
+                beta1,
+                beta2,
+                self._param_group['eps'],
+                combined_scale,
+                self._param_state['step']+1,
+                self.eps_mode,
+                bias_correction,
+                self._param_group['weight_decay'])
+
+    def _flatten_grad_mt(self, scale):
+        if self._flat_mt and len(self._grads) > 0:
+            self._overflow_buf.zero_()
+            multi_tensor_applier(
+                    amp_C.multi_tensor_scale,
+                    self._overflow_buf,
+                    list(zip(*self._grads)),
+                    scale)
+            self._grads = []
+
+    def _do_overlapped_reduction(self, param_i, param_grads_size, param_offset, param):
+        # handle overlapped reductions
+        if self._flat_mt:
+            self._grads.append( (param.grad, self._individual_flat_grads[param_i]) )
+        else:
+            torch.div(param.grad, self._world_size if self._predivide else 1.0, out=self._individual_flat_grads[param_i])
+        self._grads_generated[param_i]=True
+        if not self._last_step and self._overlap_reductions:
+            flush_block = self._get_flush_block()
+            while flush_block:
+                block_id = flush_block[0] // self._block_size
+                self._dwu_st.wait_stream(torch.cuda.current_stream())
+                with torch.cuda.stream(self._dwu_st):
+                    self._flatten_grad_mt(1.0/self._world_size if self._predivide else 1.0)
+                    torch.distributed.all_reduce(self._flat_grads_blocks[block_id])
+                if block_id == 0:
+                    self._l2_grad_norm_st.wait_stream(self._dwu_st)
+                    with torch.cuda.stream(self._l2_grad_norm_st):
+                        self._L2_grad_norm = self._flat_grads.norm(dtype=torch.float32, p=2).item()
+                flush_block = self._get_flush_block()
+
+    def set_global_scale(self, global_scale):
+        """Set global scale.
+        """
+        self._global_scale = global_scale
+
+    @property
+    def global_scale(self):
+        return self._global_scale
+
+    @property
+    def L2_grad_norm(self):
+        torch.cuda.current_stream().wait_stream(self._l2_grad_norm_st)
+        return self._L2_grad_norm
+
+    def complete_reductions(self):
+        """Complete reductions if full pipeline is not selected or overlap is not allowed.
+        """
+
+        if self._last_step:
+            # zero out gradients that have not been completed yet
+            for param_i, flat_grad in enumerate(self._individual_flat_grads):
+                if not self._grads_generated[param_i]:
+                    flat_grad.zero_()
+                    self._grads_generated[param_i] = True
+
+        if self._last_step or not self._overlap_reductions:
+            # nothing done so far, run full pipeline after reductions
+            self._dwu_st.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(self._dwu_st):
+                self._flatten_grad_mt(1.0/self._world_size if self._predivide else 1.0)
+                torch.distributed.all_reduce(self._flat_grads)
+            self._l2_grad_norm_st.wait_stream(self._dwu_st)
+            with torch.cuda.stream(self._l2_grad_norm_st):
+                self._L2_grad_norm = self._flat_grads.norm(dtype=torch.float32, p=2).item()
+
+        self._current_block = self._num_blocks
+        self._grads_generated = [False]*len(self._grads_info)
+
+    def step(self, closure=None, skip_overflow_check=False):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        with torch.cuda.stream(self._dwu_st):
+            self.__launch_step_kernel(
+                self._fp32_p,
+                self._flat_params_shards[self._rank_in_group],
+                self._fp32_m,
+                self._fp32_v,
+                self._flat_grads_shards[self._rank_in_group])
+            torch.distributed.all_gather(self._flat_params_shards, self._flat_params_shards[self._rank_in_group], group=self._ag_pg, no_copy=True)
+            for p in self._model_params: self.state[p]['step'] += 1
+
+        torch.cuda.current_stream().wait_stream(self._dwu_st)
+
+        return loss
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_lamb.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_lamb.py
new file mode 100644
index 0000000000..75bd861670
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/distributed_fused_lamb.py
@@ -0,0 +1,975 @@
+import os
+import math
+import torch
+import importlib
+import amp_C
+from apex.multi_tensor_apply import multi_tensor_applier
+
+import torch.distributed.distributed_c10d as c10d
+
+class DistributedFusedLAMB(torch.optim.Optimizer):
+
+    """Implements LAMB algorithm.
+    
+    Currently GPU-only.  Requires Apex to be installed via
+    ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``.
+    
+    This version of fused LAMB implements 2 fusions.
+      
+      * Fusion of the LAMB update's elementwise operations
+      * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
+    
+    :class:`apex.optimizers.FusedLAMB`'s usage is identical to any ordinary Pytorch optimizer::
+        
+        opt = apex.optimizers.FusedLAMB(model.parameters(), lr = ....)
+        ...
+        opt.step()
+    
+    :class:`apex.optimizers.FusedLAMB` may be used with or without Amp.  If you wish to use :class:`FusedLAMB` with Amp,
+    you may choose any ``opt_level``::
+        
+        opt = apex.optimizers.FusedLAMB(model.parameters(), lr = ....)
+        model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
+        ...
+        opt.step()
+    
+    In general, ``opt_level="O1"`` is recommended.
+    
+    LAMB was proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
+    
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its norm. (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            NOT SUPPORTED now! (default: False)
+        adam_w_mode (boolean, optional): Apply L2 regularization or weight decay
+            True for decoupled weight decay(also known as AdamW) (default: True)
+        grad_averaging (bool, optional): whether apply (1-beta2) to grad when
+            calculating running averages of gradient. (default: True)
+        set_grad_none (bool, optional): whether set grad to None when zero_grad()
+            method is called. (default: True)
+        max_grad_norm (float, optional): value used to clip global grad norm
+            (default: 1.0)
+        use_nvlamb (boolean, optional): Apply adaptive learning rate to 0.0
+            weight decay parameter (default: False)
+        step_supports_amp_scaling(boolean, optional): whether to use customized
+            gradient unscaling logic (default: True)
+    
+    .. _Large Batch Optimization for Deep Learning - Training BERT in 76 minutes:
+        https://arxiv.org/abs/1904.00962
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    class AtomicCounter(object):
+        def __init__(self):
+            self.value = 0
+            self.order = []
+            import threading
+            self._lock = threading.Lock()
+
+        def add(self, idx):
+            with self._lock:
+                self.value += 1
+                self.order.append(idx)
+
+    def __init__(self, params,
+                 lr=1e-3, bias_correction = True, grad_averaging=True,
+                 betas=(0.9, 0.999), eps=1e-8, 
+                 weight_decay=0., max_grad_norm=0., 
+                 adam_w_mode=True, use_nvlamb=False,
+                 step_supports_amp_scaling=True, overlap_reductions=True,
+                 dwu_group_size=0, dwu_num_blocks=4, dwu_num_chunks=4,
+                 dwu_num_rs_pg=1, dwu_num_ar_pg=4, dwu_num_ag_pg=0, fused_norm=False,
+                 e5m2_allgather=False, verbose=False, clip_after_ar=True,
+                 full_ar=False, set_param_views_to_flat_buffer=False, skip_allgather=False,
+                 fuse_scale=False, param_order=None, nccl_allgather_channels=0):
+        defaults = dict(lr=lr, bias_correction=bias_correction,
+                        betas=betas, eps=eps, weight_decay=weight_decay,
+                        grad_averaging=grad_averaging,
+                        max_grad_norm=max_grad_norm)
+
+        super(DistributedFusedLAMB, self).__init__(params, defaults)
+
+        global fused_adam_cuda, distributed_lamb_cuda
+        fused_adam_cuda = importlib.import_module("fused_adam_cuda")
+        distributed_lamb_cuda = importlib.import_module("distributed_lamb_cuda")
+
+        self._overflow_buf = torch.cuda.IntTensor([0])
+        self._has_overflow = False
+        self.multi_tensor_lamb_compute_update_term = distributed_lamb_cuda.multi_tensor_lamb_compute_update_term
+        self.multi_tensor_lamb_update_weights = distributed_lamb_cuda.multi_tensor_lamb_update_weights
+        import amp_C
+        self.multi_tensor_l2norm = amp_C.multi_tensor_l2norm
+
+        self._grad_averaging = grad_averaging
+        self._adam_w_mode = 1 if adam_w_mode else 0
+        self._use_nvlamb = use_nvlamb
+        self._step_supports_amp_scaling = step_supports_amp_scaling
+        self._is_accumulation_step = False
+        self._last_step = False
+        self._overlap_reductions = overlap_reductions
+        self._global_scale = None
+        self._num_blocks = dwu_num_blocks
+        self._num_chunks = dwu_num_chunks
+        self._e5m2_allgather = e5m2_allgather
+        self._verbose = verbose
+        self._clip_after_ar = clip_after_ar
+        self._full_ar = full_ar
+        self._fuse_scale = fuse_scale 
+        self._L2_grad_norm = None
+        self._set_flat_param_view = set_param_views_to_flat_buffer
+        self._skip_ag = skip_allgather
+        self._fused_norm = fused_norm 
+        self._current_process_group = c10d._get_default_group()
+        self._available_ranks = list(c10d._pg_group_ranks[self._current_process_group].keys())
+        self._group_size = torch.cuda.device_count() if dwu_group_size <= 0 else dwu_group_size
+        self._world_size = torch.distributed.get_world_size()
+        self._num_groups = self._world_size // self._group_size
+        self._rank_in_group = torch.distributed.get_rank() % self._group_size
+
+        self._lr = torch.tensor(0.0, dtype=torch.float32, device='cuda')
+
+        self._resume_from_checkpoint = False
+        self._step = torch.cuda.IntTensor([0])
+
+        # Master weight, moment, gradient buffers
+        self._fp32_p, self._fp32_m, self._fp32_v, self._fp16_p, self._fp16_g = None, None, None, None, None
+
+        import inspect
+        assert ('no_copy' in inspect.getfullargspec(torch.distributed.reduce_scatter).args), "This version of c10d does not support no_copy option"
+
+        self._num_rs_pg = dwu_num_rs_pg
+        self._num_ar_pg = dwu_num_ar_pg
+        self._num_ag_pg = dwu_num_ag_pg
+
+        if self._full_ar: # full all reduce, only need AR and AG groups
+            self._ar_pg = []
+            # consider all the ranks
+            ranks = list(range(0, self._world_size))
+            l2_grad_norm_pg = torch.distributed.new_group(ranks=ranks)
+            if torch.distributed.get_rank() in ranks:
+                self._l2_grad_norm_pg = l2_grad_norm_pg
+            for i in range(self._num_ar_pg):
+                if self._verbose:
+                    print(f"creating new AR group {i}: {ranks}")
+                grp = torch.distributed.new_group(ranks=ranks)
+                if grp != torch.distributed.GroupMember.NON_GROUP_MEMBER:
+                    if self._verbose:
+                        print(f"group {i}: init barrier (device: {torch.cuda.current_device()})")
+                    torch.distributed.barrier(group=grp, device_ids=[torch.cuda.current_device()])
+                if self._verbose:
+                    print(f"created new AR group {i}: {ranks}")
+
+                if torch.distributed.get_rank() in ranks:
+                    self._ar_pg.append(grp)
+            self._ar_st = [torch.cuda.Stream() for _ in range(self._num_ar_pg)]
+            if nccl_allgather_channels > 0:
+                os.putenv('NCCL_MAX_NCHANNELS', str(nccl_allgather_channels))
+            if self._num_ag_pg == 0:
+                self._ag_pg = self._ar_pg
+                self._ag_st = self._ar_st
+                self._num_ag_pg = self._num_ar_pg
+            else:
+                self._ag_pg = []
+                ranks = []
+                stride = torch.cuda.device_count()
+                for i in range(self._num_groups):
+                    rs = list(range(i*stride, (i+1)*stride))
+                    ranks.append(rs)
+                for rs in ranks:
+                    for i in range(self._num_ag_pg):
+                        grp = torch.distributed.new_group(ranks=rs)
+                        if torch.distributed.get_rank() in rs:
+                            if self._verbose:
+                                print(f"creating AG group {i}: {rs}")
+                            self._ag_pg.append(grp)
+
+                self._ag_st = [torch.cuda.Stream() for _ in range(self._num_ag_pg)]
+        else: # reduce-scatter + all-reduce, need RS, AR, AG groups
+            if self._num_groups > 1:
+                self._ar_pg = []
+                for dev_i in range(self._group_size):
+                    ranks = [dev_i+j*self._group_size for j in range(self._num_groups)]
+                    for i in range(self._num_ar_pg):
+                        if self._verbose:
+                            print(f"creating new AR group {i}: {ranks}")
+                        grp = torch.distributed.new_group(ranks=ranks)
+                        if grp != torch.distributed.GroupMember.NON_GROUP_MEMBER:
+                            if self._verbose:
+                                print(f"group {i}: init barrier (device: {torch.cuda.current_device()})")
+                            torch.distributed.barrier(group=grp, device_ids=[torch.cuda.current_device()])
+                        if self._verbose:
+                            print(f"created new AR group {i}: {ranks}")
+
+                        if torch.distributed.get_rank() in ranks:
+                            self._ar_pg.append(grp)
+                self._ar_st = [torch.cuda.Stream() for _ in range(self._num_ar_pg)]
+            rs_ranks = []
+            for group_i in range(self._num_groups):
+                rs_ranks.append([group_i*self._group_size+j for j in range(self._group_size)])
+            self._rs_pg = []
+            for group_i in range(self._num_groups):
+                ranks = rs_ranks[group_i]
+                for i in range(self._num_rs_pg):
+                    grp = torch.distributed.new_group(ranks=ranks)
+                    if torch.distributed.get_rank() in ranks:
+                        self._rs_pg.append(grp)
+                        if self._verbose:
+                            print(f"creating RS group : {ranks}")
+                l2_grad_norm_pg = torch.distributed.new_group(ranks=ranks)
+                if torch.distributed.get_rank() in ranks:
+                    self._l2_grad_norm_pg = l2_grad_norm_pg
+            self._rs_st = [torch.cuda.Stream() for _ in range(self._num_rs_pg)]
+            if self._num_ag_pg == 0:
+                self._ag_pg = self._rs_pg
+                self._ag_st = self._rs_st
+                self._num_ag_pg = self._num_rs_pg
+            else:
+                self._ag_pg = []
+                for group_i in range(self._num_groups):
+                    ranks = rs_ranks[group_i]
+                    for i in range(self._num_ag_pg):
+                        grp = torch.distributed.new_group(ranks=ranks)
+                        if torch.distributed.get_rank() in ranks:
+                            self._ag_pg.append(grp)
+                            if self._verbose:
+                                print(f"creating AG group : {ranks}")
+                self._ag_st = [torch.cuda.Stream() for _ in range(self._num_ag_pg)]
+        for ag_pg in self._ag_pg:
+            torch.distributed.barrier(group=ag_pg)
+
+        self._l2_grad_norm_st = torch.cuda.Stream()
+        self._completion_st = torch.cuda.Stream()
+        self._step.record_stream(self._completion_st)
+
+        self._reductions_works = [None]*self._num_blocks
+        self._allgather_works = [None]*self._num_blocks
+
+        self._one = torch.cuda.IntTensor([1])
+
+        self._first_step = True
+        self._lazy_init_stage1_done, self._lazy_init_stage2_done = False, False
+        self._param_order = self.AtomicCounter()
+
+        p_offset = 0
+        p_i = 0
+        self._model_params = []
+        self._grad_accs = []
+        self._group_properties = []
+        for group in self.param_groups:
+            prev = None
+            beta1, beta2 = group['betas']
+            beta3 = 1.0 - beta1 if self._grad_averaging else 1.0
+            bias_correction = 1 if group['bias_correction'] else 0
+            eps = group['eps']
+            weight_decay = group['weight_decay']
+            for p in group['params']:
+                if not p.requires_grad:
+                    continue
+                self._model_params.append(p)
+                self._group_properties.append((
+                    weight_decay,
+                    bias_correction,
+                    beta1,
+                    beta2,
+                    beta3,
+                    eps
+                    ))
+                p_grads_size = p.numel()
+                if self._set_flat_param_view:
+                    if param_order:
+                        # this is executed when param_order is specified by the user
+                        self._param_order.add(param_order[p])
+                    else:
+                        self._param_order.add(p_i)
+                p_offset += p_grads_size
+                # Only enforce 128b alignment (64 * fp16) for non-consecutive parameters
+                # RNN is one example of consecutive parameters:
+                # (weight_ih, weight_hh, bias_ih, bias_hh)
+                if prev is not None and (prev.data_ptr() + prev.numel() * prev.element_size() != p.data_ptr()):
+                    p_offset = ((p_offset + 63) // 64) * 64
+                prev = p
+                p_i += 1
+        if param_order:
+            self._param_order.order = torch.argsort(torch.tensor(self._param_order.order)).tolist()
+        self._grads_generated = [False]*len(self._model_params)
+        self._grads_fp16, self._grads_fp32 = [], []
+        if self._overlap_reductions:
+            self._current_block = self._num_blocks
+
+        self._net_total_param_size = p_offset
+        self._total_param_size = p_offset
+        dwu_min_page_size = 256 * self._num_blocks * self._num_chunks * self._group_size
+        self._total_param_size = ((self._total_param_size + dwu_min_page_size - 1) // dwu_min_page_size) * dwu_min_page_size
+        self._new_params = torch.zeros([self._total_param_size], dtype=torch.uint8 if self._e5m2_allgather else torch.float16, device='cuda')
+
+
+
+    def _lazy_init_stage1(self):
+        if self._lazy_init_stage1_done: return
+
+        p_i = 0
+        #self._model_params = []
+        #self._grad_accs = []
+        #self._group_properties = []
+        for group in self.param_groups:
+            for p in group['params']:
+                torch.distributed.broadcast(p, 0)
+                if not p.requires_grad:
+                    continue
+                def wrapper(param, param_i):
+                    param_tmp = param.expand_as(param)
+                    grad_acc = param_tmp.grad_fn.next_functions[0][0]
+                    def allreduce_hook(*unused):
+                        if not self._set_flat_param_view:
+                            if self._first_step:
+                                # first time
+                                self._param_order.add(param_i)
+                            else:
+                                idx = self._param_order.order.index(param_i)
+                                self._do_overlapped_reduction(idx, param)
+                        else:
+                            if not self._first_step:
+                                idx = self._param_order.order.index(param_i)
+                                self._do_overlapped_reduction(idx, param)
+                    grad_acc.register_hook(allreduce_hook)
+                    self._grad_accs.append(grad_acc)
+                wrapper(p, p_i)
+                p_i += 1
+
+        self._block_size = self._total_param_size // self._num_blocks
+        self._chunk_size = self._block_size // self._num_chunks
+        self._shard_size = self._chunk_size // self._group_size
+
+        self._flat_grads = torch.zeros([self._total_param_size], dtype=torch.float16, device='cuda')
+        self._mega_shard_size = self._num_blocks * self._num_chunks * self._shard_size
+        # initialize master weights, moments buffers if not loaded from checkpoint
+        if self._fp32_p is None:
+            self._fp32_p = torch.zeros([self._mega_shard_size], dtype=torch.float32, device='cuda')
+            self._fp32_m = torch.zeros([self._mega_shard_size], dtype=torch.float32, device='cuda')
+            self._fp32_v = torch.zeros([self._mega_shard_size], dtype=torch.float32, device='cuda')
+            self._fp32_u = torch.zeros([self._mega_shard_size], dtype=torch.float32, device='cuda')
+        # FIXME: Rethink fp16 label since it's either uint8 or fp16
+        self._fp16_p = torch.zeros([self._mega_shard_size], dtype=torch.uint8 if self._e5m2_allgather else torch.float16, device='cuda')
+        self._fp16_g = torch.zeros([self._mega_shard_size], dtype=torch.float16, device='cuda')
+
+        def _flat_split(p):
+            def __blockify(p):
+                return [p[block_id*self._block_size:(block_id+1)*self._block_size] for block_id in range(self._num_blocks)]
+            def __chunkify(p):
+                return [p[chunk_id*self._chunk_size:(chunk_id+1)*self._chunk_size] for chunk_id in range(self._num_chunks)]
+            def __shardify(p):
+                return [p[shard_id*self._shard_size:(shard_id+1)*self._shard_size] for shard_id in range(self._group_size)]
+            list_of_blocks = __blockify(p)
+            list_of_list_of_chunks = [__chunkify(block) for block in list_of_blocks]
+            list_of_list_of_list_of_shards = [[__shardify(chunk) for chunk in chunks] for chunks in list_of_list_of_chunks]
+            return list_of_blocks, list_of_list_of_chunks, list_of_list_of_list_of_shards
+        def _flat_split_no_shards(p):
+            def __blockify(p):
+                return [p[block_id*self._block_size:(block_id+1)*self._block_size] for block_id in range(self._num_blocks)]
+            def __chunkify(p):
+                return [p[chunk_id*self._chunk_size:(chunk_id+1)*self._chunk_size] for chunk_id in range(self._num_chunks)]
+            list_of_blocks = __blockify(self._flat_grads)
+            list_of_list_of_chunks = [__chunkify(block) for block in list_of_blocks]
+            return list_of_blocks, list_of_list_of_chunks
+        def _full_packed_split(p):
+            def __shardify(p):
+                return [p[mega_shard*self._mega_shard_size:(mega_shard+1)*self._mega_shard_size] for mega_shard in range(self._group_size)]
+            def __blockify(p):
+                return [p[block_id*self._num_chunks*self._shard_size:(block_id+1)*self._num_chunks*self._shard_size] for block_id in range(self._num_blocks)]
+            def __chunkify(p):
+                return [p[chunk_id*self._shard_size:(chunk_id+1)*self._shard_size] for chunk_id in range(self._num_chunks)]
+            list_of_mega_shards = __shardify(p)
+            list_of_list_of_mega_blocks = [__blockify(mega_shard) for mega_shard in list_of_mega_shards]
+            list_of_list_of_list_of_mega_chunks = [[__chunkify(mega_block) for mega_block in mega_blocks] for mega_blocks in list_of_list_of_mega_blocks]
+            return list_of_mega_shards, list_of_list_of_mega_blocks, list_of_list_of_list_of_mega_chunks
+        def _packed_split(p):
+            def __packed_blockify(p):
+                packed_block_size = self._num_chunks*self._shard_size
+                return [p[block_id*packed_block_size:(block_id+1)*packed_block_size] for block_id in range(self._num_blocks)]
+            def __packed_chunkify(p):
+                # in the packed format, each chunk contains one shard, so packed_chunk_size == self._shard_size
+                return [p[chunk_id*self._shard_size:(chunk_id+1)*self._shard_size] for chunk_id in range(self._num_chunks)]
+            list_of_blocks = __packed_blockify(p)
+            list_of_list_of_chunks = [__packed_chunkify(block) for block in list_of_blocks]
+            return list_of_blocks, list_of_list_of_chunks
+        def _split_assign(shards):
+            packed_block_size = self._num_chunks*self._shard_size
+            list_of_list_of_chunks=[]
+            for block_id in range(self._num_blocks):
+                list_of_chunks=[]
+                for chunk_id in range(self._num_chunks):
+                    #self._fp16_g[block_id*packed_block_size+chunk_id*self._shard_size:block_id*packed_block_size+(chunk_id+1)*self._shard_size] = shards[block_id][chunk_id][self._rank_in_group]
+                    list_of_chunks.append( shards[block_id][chunk_id][self._rank_in_group])
+                list_of_list_of_chunks.append(list_of_chunks)
+            return list_of_list_of_chunks
+
+        self._new_params_mega_shards, self._new_params_mega_blocks, self._new_params_mega_chunks = _full_packed_split(self._new_params)
+        # this splitting scheme is needed when allgather needs to be split into multiple chunks in a contiguous way
+        self._new_params2_blocks, self._new_params2_chunks, self._new_params2_shards = _flat_split(self._new_params)
+
+        self._fp32_p_blocks, self._fp32_p_chunks = _packed_split(self._fp32_p)
+        self._fp32_m_blocks, self._fp32_m_chunks = _packed_split(self._fp32_m)
+        self._fp32_v_blocks, self._fp32_v_chunks = _packed_split(self._fp32_v)
+        self._fp32_u_blocks, self._fp32_u_chunks = _packed_split(self._fp32_u)
+        self._fp16_p_blocks, self._fp16_p_chunks = _packed_split(self._fp16_p)
+
+        if self._full_ar:
+            # for gradient all-reduce
+            self._flat_grads_blocks, self._flat_grads_chunks, self._flat_grads_shards = _flat_split(self._flat_grads)
+            # for weight update
+            self._fp16_g_chunks = _split_assign(self._flat_grads_shards)
+        else:
+            self._flat_grads_blocks, self._flat_grads_chunks, self._flat_grads_shards = _flat_split(self._flat_grads)
+            self._fp16_g_blocks, self._fp16_g_chunks = _packed_split(self._fp16_g)
+
+        self._lazy_init_stage1_done = True
+
+    def _lazy_init_stage2(self):
+        if self._lazy_init_stage2_done: return
+        if not self._set_flat_param_view: 
+            # reversing is needed for overlapping allreduce and backprop, but currently not supported for flat param view
+            self._param_order.order.reverse()
+
+            # re-order model_params, grad_accs, group_properties lists
+        self._model_params = [self._model_params[i] for i in self._param_order.order]
+        self._grad_accs = [self._grad_accs[i] for i in self._param_order.order]
+        self._group_properties = [self._group_properties[i] for i in self._param_order.order]
+
+        def _get_flat_view(param):
+            if param.is_contiguous(memory_format=torch.channels_last):
+                K, C, H, W = param.shape
+                pv = param.as_strided(size=(K,H,W,C), stride=(H*W*C, W*C, C, 1))
+            elif param.is_contiguous(memory_format=torch.channels_last_3d):
+                K, C, D, H, W = param.shape
+                pv = param.as_strided(size=(K,D,H,W,C), stride=(D*H*W*C, H*W*C, W*C, C, 1))
+            else:
+                pv = param
+            return pv.view(-1)
+
+        # re-collect grads info (size, offset) after ordering
+        prev = None
+        p_offset = 0
+        self._grads_info = []
+        self._individual_flat_grads = []
+        for i, p in enumerate(self._model_params):
+            p_grads_size = p.numel()
+            self._grads_info.append({"param_grads_size":p_grads_size, "param_offset":p_offset})
+            self._individual_flat_grads.append(self._flat_grads[p_offset:p_offset+p_grads_size].view_as(p))
+            # for the first iteration
+            self._do_overlapped_reduction(i, p)
+            p_offset += p_grads_size
+            # Only enforce 128b alignment (64 * fp16) for non-consecutive parameters
+            # RNN is one example of consecutive parameters:
+            # (weight_ih, weight_hh, bias_ih, bias_hh)
+            if prev is not None and (prev.data_ptr() + prev.numel() * prev.element_size() != p.data_ptr()):
+                p_offset = ((p_offset + 63) // 64) * 64
+            prev = p
+
+        self._low_param_i = [0]*self._num_blocks
+        for block_id in range(self._num_blocks-1,-1,-1):
+            p_i = len(self._grads_info)-1
+            while p_i > 0 and self._grads_info[p_i]["param_offset"] > block_id*self._block_size:
+                p_i -= 1
+            self._low_param_i[block_id] = p_i
+        #print("self._low_param_i", self._low_param_i)
+
+        # This paragraph does two things:
+        # 1) Copy model parameters into master buffer
+        # 2) Create tensor lists for unpacking new parameter tensor after all-gather
+        self._packed_flat_to_model_params_fp16 = []
+        self._packed_flat_to_model_params_fp32 = []
+        self._model_params_num = len(self._model_params)
+        self._contrib_tensor_list = []
+        self._contrib_min_param_i, self._contrib_max_param_i = -1, -1
+        self._contrib_update_frag_for_norm = []
+        self._contrib_model_param_for_norm_fp16 = []
+        self._contrib_model_param_for_norm_fp32 = []
+        self._contrib_model_param_for_norm_is_fp16 = []
+        self._model_param_is_contrib = []
+        self._contrib_group_properties = []
+        for shard_id in range(self._group_size):
+            for block_id in range(self._num_blocks):
+                for chunk_id in range(self._num_chunks):
+                    flat_shard_start = (((block_id * self._num_chunks + chunk_id) * self._group_size) + shard_id) * self._shard_size
+                    flat_shard_end = flat_shard_start + self._shard_size
+                    for param_i, (p, grads_info, group_props) in enumerate(zip(self._model_params, self._grads_info, self._group_properties)):
+                        flat_grad_start = grads_info["param_offset"]
+                        flat_grad_end = flat_grad_start + grads_info["param_grads_size"]
+                        clipped_start = (lambda a,b: a if a > b else b)(flat_grad_start, flat_shard_start)
+                        clipped_end = (lambda a,b: a if a < b else b)(flat_grad_end, flat_shard_end)
+                        if clipped_start < clipped_end:
+                            grad_offset = clipped_start - flat_grad_start
+                            grad_length = clipped_end - clipped_start
+                            shard_offset = clipped_start - flat_shard_start
+                            pf = _get_flat_view(p)
+                            model_param_fragment = pf[grad_offset:grad_offset+grad_length]
+                            new_param_packed_fragment = self._new_params_mega_chunks[shard_id][block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                            if model_param_fragment.dtype == torch.float16:
+                                self._packed_flat_to_model_params_fp16.append( (new_param_packed_fragment, model_param_fragment) )
+                            else:
+                                self._packed_flat_to_model_params_fp32.append( (new_param_packed_fragment, model_param_fragment) )
+                            if shard_id == self._rank_in_group:
+                                self._model_param_is_contrib.append(param_i)
+                                # copy model parameters into master buffer
+                                master_param_fragment = self._fp32_p_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                opti_state_m_fragment = self._fp32_m_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                opti_state_v_fragment = self._fp32_v_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                opti_state_u_fragment = self._fp32_u_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                opti_state_g_fragment = self._fp16_g_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                opti_state_p_fragment = self._fp16_p_chunks[block_id][chunk_id][shard_offset:shard_offset+grad_length]
+                                #print("model_param_fragment.size()=%s, new_param_packed_fragment.size()=%s, master_param_fragment.size()=%s" % (str(model_param_fragment.size()), str(new_param_packed_fragment.size()), str(master_param_fragment.size())))
+                                if not self._resume_from_checkpoint:
+                                    master_param_fragment.copy_(model_param_fragment)
+                                self._contrib_group_properties.append(group_props)
+                                self._contrib_tensor_list.append((master_param_fragment, opti_state_m_fragment, opti_state_v_fragment, opti_state_u_fragment, opti_state_g_fragment, opti_state_p_fragment)) # p, m, v, u, g, p_copy
+                                self._contrib_update_frag_for_norm.append(opti_state_u_fragment)
+                                if p.dtype == torch.float16:
+                                    self._contrib_model_param_for_norm_fp16.append(p)
+                                else:
+                                    self._contrib_model_param_for_norm_fp32.append(p)
+                                self._contrib_model_param_for_norm_is_fp16.append(True if p.dtype == torch.float16 else False)
+                                if self._contrib_min_param_i < 0: self._contrib_min_param_i = param_i
+                                self._contrib_max_param_i = param_i
+        self._contrib_model_param_for_norm_num = len(self._contrib_model_param_for_norm_is_fp16)
+        if len(self._contrib_model_param_for_norm_fp16) == 0: self._contrib_model_param_for_norm_fp16 = None
+        if len(self._contrib_model_param_for_norm_fp32) == 0: self._contrib_model_param_for_norm_fp32 = None
+        self._contrib_model_param_for_norm_is_fp32 = torch.tensor([not is_fp16 for is_fp16 in self._contrib_model_param_for_norm_is_fp16], dtype=torch.bool, device='cuda')
+        self._contrib_model_param_for_norm_is_fp16 = torch.tensor([is_fp16 for is_fp16 in self._contrib_model_param_for_norm_is_fp16], dtype=torch.bool, device='cuda')
+        self._offsets = torch.tensor(self._model_param_is_contrib, dtype=torch.int64, device='cuda')
+
+        p, m, v, u, g, p_copy = list(zip(*self._contrib_tensor_list))
+        self._contrib_compute_update_term_tensor_list = [g, p, m, v, u]
+        self._contrib_update_weights_tensor_list = [u, p, p_copy]
+
+        math_type = self._fp32_u.dtype
+        decay, bias_correction, beta1, beta2, beta3, epsilon = list(zip(*self._contrib_group_properties))
+        self._contrib_beta1 = torch.tensor(beta1, dtype=math_type, device='cuda')
+        self._contrib_beta2 = torch.tensor(beta2, dtype=math_type, device='cuda')
+        self._contrib_beta3 = torch.tensor(beta3, dtype=math_type, device='cuda')
+        self._contrib_bias_correction = torch.tensor(bias_correction, dtype=torch.int, device='cuda')
+        self._contrib_epsilon = torch.tensor(epsilon, dtype=math_type, device='cuda')
+        self._contrib_weight_decay = torch.tensor(decay, dtype=math_type, device='cuda')
+
+        self._packed_flat_to_model_params_fp16 = list(zip(*self._packed_flat_to_model_params_fp16)) if len(self._packed_flat_to_model_params_fp16) > 0 else None
+        self._packed_flat_to_model_params_fp32 = list(zip(*self._packed_flat_to_model_params_fp32)) if len(self._packed_flat_to_model_params_fp32) > 0 else None
+
+        self._lazy_init_stage2_done = True
+
+        self.complete_reductions()
+        self._first_step = False
+
+    def set_is_accumulation_step(self, is_accumulation_step):
+        self._is_accumulation_step = is_accumulation_step
+
+    def set_last_step(self, last_step):
+        self._last_step = last_step
+        
+    def _get_flush_block(self):
+        flush_block = []
+        if self._current_block > 0 and self._grads_generated[self._low_param_i[self._current_block-1]]:
+            num_grads = len(self._grads_generated)
+            contiguous_idx = num_grads
+            while contiguous_idx > 0 and self._grads_generated[contiguous_idx-1]:
+                contiguous_idx -= 1
+
+            if contiguous_idx < num_grads and self._grads_info[contiguous_idx]["param_offset"] <= (self._current_block-1)*self._block_size:
+                self._current_block -= 1
+                start = self._current_block * self._block_size
+                end = (self._current_block+1) * self._block_size
+                flush_block = [start, end]
+
+        return flush_block
+
+    def _full_all_reduce_scale(self, block_id, scale):
+        works = [None]*self._num_chunks
+        if  self._clip_after_ar:    
+            for chunk_id in range(self._num_chunks):
+                glob_chunk_id = block_id * self._num_chunks + chunk_id
+                ar_stream = self._ar_st[glob_chunk_id%self._num_ar_pg]
+                ar_stream.wait_stream(torch.cuda.current_stream())
+                with torch.cuda.stream(ar_stream):
+                    works[chunk_id] = torch.distributed.all_reduce(self._flat_grads_chunks[block_id][chunk_id],group=self._ar_pg[glob_chunk_id%self._num_ar_pg],async_op=True,op=torch.distributed.make_nccl_premul_sum((scale,)))
+        else:
+            glob_chunk_id = block_id
+            ar_stream = self._ar_st[glob_chunk_id%self._num_ar_pg]
+            ar_stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(ar_stream):
+                    works0 = torch.distributed.all_reduce(self._flat_grads_blocks[block_id],group=self._ar_pg[glob_chunk_id%self._num_ar_pg],async_op=True,op=torch.distributed.make_nccl_premul_sum((scale,)))
+            for i in range(self._num_chunks):
+                works[i]=works0
+        self._reductions_works[block_id] = works
+
+    def _full_all_reduce(self, block_id):
+        works = [None]*self._num_chunks
+
+        for chunk_id in range(self._num_chunks):
+            glob_chunk_id = block_id * self._num_chunks + chunk_id
+            ar_stream = self._ar_st[glob_chunk_id%self._num_ar_pg]
+            ar_stream.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(ar_stream):
+                works[chunk_id] = torch.distributed.all_reduce(self._flat_grads_chunks[block_id][chunk_id],group=self._ar_pg[glob_chunk_id%self._num_ar_pg],async_op=True)
+        self._reductions_works[block_id] = works
+
+    def _reduce_scatter_and_all_reduce_scale(self, block_id, scale):
+        # Reduction within each node
+        # Changes gradient format from [block * chunk * shard] to [shard * block * chunk]
+        # The output format is the same as the fp32 master parameters
+        works = [None]*self._num_chunks
+        for chunk_id in range(self._num_chunks):
+            glob_chunk_id = block_id * self._num_chunks + chunk_id
+            rs_stream = self._rs_st[glob_chunk_id%self._num_rs_pg]
+            rs_stream.wait_stream(torch.cuda.current_stream())
+            rs_stream.wait_stream(self._l2_grad_norm_st)
+            with torch.cuda.stream(rs_stream):
+                works[chunk_id] = torch.distributed.reduce_scatter(self._fp16_g_chunks[block_id][chunk_id],self._flat_grads_shards[block_id][chunk_id],group=self._rs_pg[glob_chunk_id%self._num_rs_pg],async_op=True,no_copy=True,op=torch.distributed.make_nccl_premul_sum((scale,)))
+
+        # Reduction across nodes for each rank
+        if self._num_groups > 1:
+            for chunk_id in range(self._num_chunks):
+                glob_chunk_id = block_id * self._num_chunks + chunk_id
+                ar_stream = self._ar_st[glob_chunk_id%self._num_ar_pg]
+                with torch.cuda.stream(ar_stream):
+                    works[chunk_id].wait()
+                    works[chunk_id] = torch.distributed.all_reduce(self._fp16_g_chunks[block_id][chunk_id],group=self._ar_pg[glob_chunk_id%self._num_ar_pg],async_op=True)
+        self._reductions_works[block_id] = works
+
+    def _reduce_scatter_and_all_reduce(self, block_id):
+        # Reduction within each node
+        # Changes gradient format from [block * chunk * shard] to [shard * block * chunk]
+        # The output format is the same as the fp32 master parameters
+        works = [None]*self._num_chunks
+        for chunk_id in range(self._num_chunks):
+            glob_chunk_id = block_id * self._num_chunks + chunk_id
+            rs_stream = self._rs_st[glob_chunk_id%self._num_rs_pg]
+            rs_stream.wait_stream(torch.cuda.current_stream())
+            rs_stream.wait_stream(self._l2_grad_norm_st)
+            with torch.cuda.stream(rs_stream):
+                works[chunk_id] = torch.distributed.reduce_scatter(self._fp16_g_chunks[block_id][chunk_id],self._flat_grads_shards[block_id][chunk_id],group=self._rs_pg[glob_chunk_id%self._num_rs_pg],async_op=True,no_copy=True)
+
+        # Reduction across nodes for each rank
+        if self._num_groups > 1:
+            for chunk_id in range(self._num_chunks):
+                glob_chunk_id = block_id * self._num_chunks + chunk_id
+                ar_stream = self._ar_st[glob_chunk_id%self._num_ar_pg]
+                with torch.cuda.stream(ar_stream):
+                    works[chunk_id].wait()
+                    works[chunk_id] = torch.distributed.all_reduce(self._fp16_g_chunks[block_id][chunk_id],group=self._ar_pg[glob_chunk_id%self._num_ar_pg],async_op=True)
+        self._reductions_works[block_id] = works
+
+    def _pipeline_block_reductions(self, block_id):
+        if self._clip_after_ar:
+            self._flatten_grad_mt(1.0/self._world_size)
+
+            if self._full_ar:
+                self._full_all_reduce(block_id)
+            else:
+                self._reduce_scatter_and_all_reduce(block_id)
+
+            # Compute L2 grad norm
+            if block_id == 0:
+                with torch.cuda.stream(self._l2_grad_norm_st):
+                    for block_id in range(self._num_blocks):
+                        for chunk_id in range(self._num_chunks):
+                            self._reductions_works[block_id][chunk_id].wait()
+                    # Since the packed format is contiguous after reductions, only one norm is needed
+                    l2_grad_norm_sq = torch.empty([1], device='cuda')
+                    if 0:#self._full_ar:
+                        l2_grad_norm_sq = self._flat_grads_shards[self._rank_in_group].norm(dtype=torch.float32, p=2)**2
+                    else:
+                        l2_grad_norm_sq = self._fp16_g.norm(dtype=torch.float32, p=2)**2
+                    torch.distributed.all_reduce(l2_grad_norm_sq, group=self._l2_grad_norm_pg)
+                    self._L2_grad_norm = l2_grad_norm_sq.sqrt()
+        else:
+            # Copy model grads to flat grads buffer
+            self._flatten_grad_mt(1.0)
+
+            # Compute L2 grad norm
+            self._l2_grad_norm_st.wait_stream(torch.cuda.current_stream())
+            with torch.cuda.stream(self._l2_grad_norm_st):
+                if not self._fused_norm:
+                    self._L2_grad_norm = self._flat_grads.norm(dtype=torch.float16, p=2).float()
+            torch.cuda.current_stream().wait_stream(self._l2_grad_norm_st)
+
+            # Apply clipping & pre-reduction scaling on grads
+            loss_scale = self.global_scale
+            max_grad_norm = loss_scale*self.defaults['max_grad_norm']
+            coeff = max_grad_norm /(1e-6+self.L2_grad_norm)
+            coeff = (coeff>1) * self._one + (coeff<=1) * coeff
+            tmp = torch.cat(((self._one), (coeff)))
+            index = (coeff+1>coeff).int()
+            scale = tmp.index_select(0, index).half()/self._world_size
+            if not self._fuse_scale:
+                self._flat_grads.mul_(scale)
+
+            if self._full_ar:
+                if self._fuse_scale:
+                    self._full_all_reduce_scale(block_id, scale)
+                else:
+                    self._full_all_reduce(block_id) 
+            else:
+                if self._fuse_scale:
+                    self._reduce_scatter_and_all_reduce_scale(block_id, scale)
+                else:
+                    self._reduce_scatter_and_all_reduce(block_id)
+
+            if block_id == 0:
+                for block_id in range(self._num_blocks):
+                    for chunk_id in range(self._num_chunks):
+                        self._reductions_works[block_id][chunk_id].wait()
+
+    def __compute_contrib_param_norm(self):
+        if self._contrib_model_param_for_norm_fp16 is not None and self._contrib_model_param_for_norm_fp32 is not None:
+            gnorm_fp16 = multi_tensor_applier(self.multi_tensor_l2norm, self._overflow_buf, [self._contrib_model_param_for_norm_fp16], True)[1]
+            gnorm_fp32 = multi_tensor_applier(self.multi_tensor_l2norm, self._overflow_buf, [self._contrib_model_param_for_norm_fp32], True)[1]
+            gnorm = torch.empty(size=[self._contrib_model_param_for_norm_num], dtype=torch.bool, device='cuda')
+            gnorm.masked_scatter_(self._contrib_model_param_for_norm_is_fp16, gnorm_fp16)
+            gnorm.masked_scatter_(self._contrib_model_param_for_norm_is_fp32, gnorm_fp32)
+        elif self._contrib_model_param_for_norm_fp16 is not None:
+            gnorm = multi_tensor_applier(self.multi_tensor_l2norm, self._overflow_buf, [self._contrib_model_param_for_norm_fp16], True)[1]
+        elif self._contrib_model_param_for_norm_fp32 is not None:
+            gnorm = multi_tensor_applier(self.multi_tensor_l2norm, self._overflow_buf, [self._contrib_model_param_for_norm_fp32], True)[1]
+        return gnorm
+
+    def __compute_contrib_update_norm(self):
+        l2_norm = torch.zeros(size=[self._model_params_num], dtype=torch.float32, device='cuda')
+        local_contrib_l2_norm = multi_tensor_applier(self.multi_tensor_l2norm, self._overflow_buf, [self._contrib_update_frag_for_norm], True)[1] ** 2
+        l2_norm.scatter_(dim=0, index=self._offsets, src=local_contrib_l2_norm)
+        torch.distributed.all_reduce(l2_norm, group=self._ag_pg[0])
+        l2_norm = torch.sqrt(l2_norm)
+        return l2_norm
+
+    def _pipeline_step(self):
+        global_scale = self.global_scale
+        # if clip before ar, set max_grad_norm to 0
+        max_grad_norm = self.defaults['max_grad_norm'] * self._clip_after_ar
+        self._completion_st.wait_stream(self._l2_grad_norm_st)
+        global_grad_norm = self.L2_grad_norm
+
+        # check global_grad_norm and fill overflow_buf
+        is_finite = (global_grad_norm + 1 > global_grad_norm).int()
+        self._overflow_buf = self._one * (is_finite ^ self._one) # toggle between 0 and 1
+
+        if not self._clip_after_ar:
+            torch.distributed.all_reduce(is_finite,
+                                         op=torch.distributed.ReduceOp.MIN,
+                                         group=self._current_process_group)
+            torch.distributed.all_reduce(self._overflow_buf,
+                                         op=torch.distributed.ReduceOp.MAX,
+                                         group=self._current_process_group)
+
+        # increment step counter if no overflow
+        self._step += is_finite
+        self._completion_st.wait_stream(torch.cuda.current_stream())
+        self._completion_st.wait_stream(self._l2_grad_norm_st)
+
+        # Call step kernel once per step
+        # Call all-gather once per step
+        with torch.cuda.stream(self._completion_st):
+            for block_id in range(self._num_blocks):
+                for chunk_id in range(self._num_chunks):
+                    self._reductions_works[block_id][chunk_id].wait()
+            param_norm = self.__compute_contrib_param_norm()
+            multi_tensor_applier(self.multi_tensor_lamb_compute_update_term,
+                    self._overflow_buf,
+                    self._contrib_compute_update_term_tensor_list, # g, p, m, v, u
+                    self._contrib_beta1,
+                    self._contrib_beta2,
+                    self._contrib_beta3,
+                    self._contrib_bias_correction,
+                    self._step,
+                    self._contrib_epsilon,
+                    self._adam_w_mode,
+                    self._contrib_weight_decay,
+                    global_scale,
+                    global_grad_norm,
+                    max_grad_norm)
+            upd_norm = self.__compute_contrib_update_norm()
+            multi_tensor_applier(self.multi_tensor_lamb_update_weights,
+                    self._overflow_buf,
+                    self._contrib_update_weights_tensor_list, # u, p, p_copy
+                    param_norm,
+                    upd_norm,
+                    self._offsets,
+                    self._lr,
+                    self._contrib_weight_decay,
+                    global_grad_norm,
+                    self._use_nvlamb)
+            if not self._skip_ag:
+                # allgather chunking is currently not supported for clip after allreduce 
+                if not self._clip_after_ar:
+                    for block in range(self._num_blocks):
+                        for chunk in range(self._num_chunks):
+                            torch.distributed.all_gather(self._new_params2_shards[block][chunk], self._fp16_p_chunks[block][chunk], group=self._ag_pg[0], no_copy=True)
+                else:
+                    torch.distributed.all_gather(self._new_params_mega_shards, self._fp16_p, group=self._ag_pg[0], no_copy=True)
+
+    def _flatten_grad_mt(self, scale):
+        if len(self._grads_fp16) > 0:
+            self._overflow_buf.zero_()
+            if not self._fused_norm:
+                multi_tensor_applier(
+                        amp_C.multi_tensor_scale,
+                        self._overflow_buf,
+                        list(zip(*self._grads_fp16)),
+                        scale)
+            else:
+                self._L2_grad_norm=multi_tensor_applier(
+                        amp_C.multi_tensor_l2norm_scale,
+                        self._overflow_buf,
+                        list(zip(*self._grads_fp16)),
+                        scale, False)[0].float()
+
+            self._grads_fp16 = []
+        if len(self._grads_fp32) > 0:
+            self._overflow_buf.zero_()
+            if not self._fused_norm:
+                multi_tensor_applier(
+                        amp_C.multi_tensor_scale,
+                        self._overflow_buf,
+                        list(zip(*self._grads_fp32)),
+                        scale)
+            else:
+                self._L2_grad_norm=multi_tensor_applier(
+                        amp_C.multi_tensor_l2norm_scale,
+                        self._overflow_buf,
+                        list(zip(*self._grads_fp32)),
+                        scale, False)[0].float()
+            self._grads_fp32 = []
+
+    def _do_overlapped_reduction(self, param_i, param):
+        if not self._is_accumulation_step:
+            # handle overlapped reductions
+            if param.dtype == torch.float16:
+                self._grads_fp16.append( (param.grad, self._individual_flat_grads[param_i]) )
+            else:
+                self._grads_fp32.append( (param.grad, self._individual_flat_grads[param_i]) )
+            self._grads_generated[param_i]=True
+            if not self._first_step and not self._last_step:
+                if self._overlap_reductions:
+                    flush_block = self._get_flush_block()
+                    while flush_block:
+                        block_id = flush_block[0] // self._block_size
+                        self._pipeline_block_reductions(block_id)
+                        flush_block = self._get_flush_block()
+
+    def set_global_scale(self, global_scale):
+        """Set global scale.
+        """
+        self._global_scale = global_scale
+
+    @property
+    def global_scale(self):
+        return self._global_scale
+
+    @property
+    def L2_grad_norm(self):
+        torch.cuda.current_stream().wait_stream(self._l2_grad_norm_st)
+        return self._L2_grad_norm
+
+    def complete_reductions(self):
+        """Complete reductions if full pipeline is not selected or overlap is not allowed.
+        """
+        if self._last_step:
+            # zero out gradients that have not been completed yet
+            for param_i, grad_generated in enumerate(self._grads_generated):
+                if not grad_generated:
+                    grad_info = self._grads_info[param_i]
+                    param_offset = grad_info["param_offset"]
+                    param_size = grad_info["param_grads_size"]
+                    self._flat_grads[param_offset:param_offset+param_size].zero_()
+                    self._grads_generated[param_i] = True
+
+        if self._first_step or self._last_step or not self._overlap_reductions:
+            # nothing done so far, run full pipeline after reductions
+            for block_id in range(self._num_blocks-1,-1,-1):
+                self._pipeline_block_reductions(block_id)
+
+        torch.cuda.current_stream().wait_stream(self._l2_grad_norm_st)
+
+        self._current_block = self._num_blocks
+        self._grads_generated = [False]*len(self._grads_info)
+
+    def step(self, closure=None, grad_scaler=None):
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        self._pipeline_step()
+
+        if grad_scaler is not None:
+            found_inf = self._overflow_buf.float()
+            optimizer_state = grad_scaler._per_optimizer_states[id(self)]
+            current_device = torch.device('cuda', torch.cuda.current_device())
+            optimizer_state["found_inf_per_device"][current_device] = found_inf
+
+        self._completion_st.wait_stream(torch.cuda.current_stream())
+        if not self._set_flat_param_view:
+            with torch.cuda.stream(self._completion_st):
+                # Copy self._new_params to model params
+                with torch.no_grad():
+                    if self._packed_flat_to_model_params_fp16 is not None:
+                        multi_tensor_applier(
+                                fused_adam_cuda.maybe_cast_mt,
+                                self._overflow_buf,
+                                self._packed_flat_to_model_params_fp16)
+                    if self._packed_flat_to_model_params_fp32 is not None:
+                        multi_tensor_applier(
+                                fused_adam_cuda.maybe_cast_mt,
+                                self._overflow_buf,
+                                self._packed_flat_to_model_params_fp32)
+    
+        torch.cuda.current_stream().wait_stream(self._completion_st)
+
+        self._reductions_works = [None]*self._num_blocks
+        self._allgather_works = [None]*self._num_blocks
+
+        return loss
+
+    def state_dict(self):
+        """
+        Returns a dict containing the current state of this :class:`DistributedFusedAdam` instance.
+        Example::
+            checkpoint = {}
+            checkpoint['model'] = model.state_dict()
+            checkpoint['optimizer'] = optimizer.state_dict()
+            torch.save(checkpoint, "saved.pth")
+        """
+        # save step, master weights and first/second moments
+        state_dict = {}
+        state_dict['step'] = self._step
+        state_dict['fp32_p'] = self._fp32_p
+        state_dict['fp32_m'] = self._fp32_m
+        state_dict['fp32_v'] = self._fp32_v
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+        """
+        Loads a state_dict created by an earlier call to state_dict().
+        If an DistributedFusedAdam instance was constructed from some ``init_optimizer``,
+        whose parameters in turn came from ``model``, it is expected that the user
+        will call ``model.load_state_dict()`` before
+        ``optimizer.load_state_dict()`` is called.
+        Example::
+            model = torch.nn.Linear(D_in, D_out).cuda().half()
+            optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
+            optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
+            ...
+            checkpoint = torch.load("saved.pth")
+            model.load_state_dict(checkpoint['model'])
+            optimizer.load_state_dict(checkpoint['optimizer'])
+        """
+        # restore step, master weights and first/second moments
+        self._step = state_dict['step']
+        self._fp32_p = state_dict['fp32_p'].to(device="cuda")
+        self._fp32_m = state_dict['fp32_m'].to(device="cuda")
+        self._fp32_v = state_dict['fp32_v'].to(device="cuda")
+        self._resume_from_checkpoint = True
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fp16_optimizer.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fp16_optimizer.py
new file mode 100644
index 0000000000..0cbb63b829
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fp16_optimizer.py
@@ -0,0 +1,243 @@
+import torch
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class FP16_Optimizer(object):
+    """
+    :class:`FP16_Optimizer` A cutdown version of apex.fp16_utils.FP16_Optimizer.
+    Designed only to wrap apex.contrib.optimizers.FusedAdam, FusedSGD.
+    Refer to apex.fp16_utils documents for more information.
+    Example::
+        model = torch.nn.Linear(D_in, D_out).cuda().half()
+        optimizer = apex.contrib.optimizers.FusedSGD(model.parameters())
+        optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
+        ...
+        # loss.backward() becomes:
+        optimizer.backward(loss)
+        ...
+    Example with dynamic loss scaling::
+        ...
+        optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
+                                   # optional arg to control dynamic loss scaling behavior
+                                   # dynamic_loss_args={'scale_window' : 500})
+                                   # Usually, dynamic_loss_args is not necessary.
+    """
+
+    def __init__(self,
+                 init_optimizer,
+                 static_loss_scale=1.0,
+                 dynamic_loss_scale=False,
+                 dynamic_loss_args=None,
+                 verbose=True):
+
+        print("\nThis fp16_optimizer is designed to only work with apex.contrib.optimizers.*")
+        print("To update, use updated optimizers with AMP.")
+        # The fused optimizer does all the work. We need this layer for two reason:
+        # 1. maintain same user API from apex.fp16_utils
+        # 2. keep common stuff here in case we need to add new fused optimizer later
+
+        if not torch.cuda.is_available:
+            raise SystemError("Cannot use fp16 without CUDA.")
+        self.optimizer = init_optimizer
+
+        self.fp16_groups = [] # model params
+        self.fp32_groups = [] # master weights
+
+        # iterate over param_groups
+        for param_group in self.optimizer.param_groups:
+            fp16_group = []
+            fp32_group = []
+            for p in param_group['params']:
+                fp16_group.append(p)
+                fp32_group.append(p.clone().float().detach())
+            self.fp16_groups.append(fp16_group)
+            self.fp32_groups.append(fp32_group)
+            param_group['params'] = fp32_group
+
+        if multi_tensor_applier.available:
+            import amp_C
+            self.overflow_buf = torch.cuda.IntTensor([0])
+            self.multi_tensor_l2norm=amp_C.multi_tensor_l2norm
+        else:
+            raise RuntimeError('FP16_Optimizer requires cuda extensions')
+
+        # we may have a way of fusing dynamic scale. Do not support for now
+        if dynamic_loss_scale:
+            if dynamic_loss_args is not None:
+                raise SystemError("Do not support dynamic loss scale args for now.")
+            self.dynamic_loss_scale = True
+            self.cur_scale = 2**16
+            self.cur_iter = 0
+            self.last_overflow_iter = -1
+            self.scale_factor = 2
+            self.scale_window = 1000
+        else:
+            self.dynamic_loss_scale = False
+            self.cur_iter = 0
+            self.cur_scale = static_loss_scale
+        self.verbose = verbose
+
+    def zero_grad(self, set_grads_to_None=True):
+        """
+        Zero FP16 parameter grads.
+        """
+        # FP32 grad should never exist.
+        # For speed, set model fp16 grad to None by default
+        for group in self.fp16_groups:
+            for p in group:
+                if set_grads_to_None:
+                    p.grad = None
+                else:
+                    if p.grad is not None:
+                        p.grad.detach_()
+                        p.grad.zero_()
+
+    def step(self, closure=None):
+        """
+        Not supporting closure.
+        """
+        fp16_grads = []
+        norm_groups = []
+        skip = False
+
+        for group in self.fp16_groups:
+            fp16_grad = []
+            for i, p in enumerate(group):
+                fp16_grad.append(p.grad)
+            fp16_grads.append(fp16_grad)
+        
+        # nan check
+        self.overflow_buf.zero_()
+        for fp16_grad in fp16_grads:
+            if len(fp16_grad) > 0:
+                norm, norm_per_tensor = multi_tensor_applier(self.multi_tensor_l2norm,
+                                                             self.overflow_buf,
+                                                             [fp16_grad], True)
+                norm_groups.append(norm)
+                if self.overflow_buf.item() != 0:
+                    skip = True
+
+        if skip:
+            self._update_scale(skip)
+            return
+
+        # norm is in fact norm*cur_scale
+        self.optimizer.step(grads=fp16_grads,
+                            output_params=self.fp16_groups,
+                            scale=self.cur_scale,
+                            grad_norms=norm_groups)
+
+        self._update_scale(False)
+        return
+
+    def backward(self, loss):
+        """
+        :attr:`backward` performs the following steps:
+        1. fp32_loss = loss.float()
+        2. scaled_loss = fp32_loss*loss_scale
+        3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's fp16 leaves
+        """
+        scaled_loss = (loss.float()) * self.cur_scale
+        scaled_loss.backward()
+
+    def _update_scale(self, skip):
+        if self.dynamic_loss_scale:
+            if skip:
+                if self.verbose:
+                    print("\nGrad overflow on iteration", self.cur_iter)
+                    print("Using dynamic loss scale of", self.cur_scale)
+                self.cur_scale = max(self.cur_scale/self.scale_factor, 1)
+                self.last_overflow_iter = self.cur_iter
+            else:
+                if (self.cur_iter - self.last_overflow_iter) % self.scale_window == 0:
+                    self.cur_scale *= self.scale_factor
+        else:
+            if skip:
+                print("\nGrad overflow on iteration", self.cur_iter)
+                print("Using static loss scale of", self.cur_scale)
+        self.cur_iter +=1
+        return
+
+    # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state"
+    def _get_state(self):
+        return self.optimizer.state
+
+    def _set_state(self, value):
+        self.optimizer.state = value
+
+    state = property(_get_state, _set_state)
+
+    # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups"
+    # (for example, to adjust the learning rate)
+    def _get_param_groups(self):
+        return self.optimizer.param_groups
+
+    def _set_param_groups(self, value):
+        self.optimizer.param_groups = value
+
+    param_groups = property(_get_param_groups, _set_param_groups)
+
+    def state_dict(self):
+        """
+        Returns a dict containing the current state of this :class:`FP16_Optimizer` instance.
+        This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict
+        of the contained Pytorch optimizer.
+        Example::
+            checkpoint = {}
+            checkpoint['model'] = model.state_dict()
+            checkpoint['optimizer'] = optimizer.state_dict()
+            torch.save(checkpoint, "saved.pth")
+        """
+        state_dict = {}
+        state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale
+        state_dict['cur_scale'] = self.cur_scale
+        state_dict['cur_iter'] = self.cur_iter
+        if state_dict['dynamic_loss_scale']:
+            state_dict['last_overflow_iter'] = self.last_overflow_iter
+            state_dict['scale_factor'] = self.scale_factor
+            state_dict['scale_window'] = self.scale_window
+        state_dict['optimizer_state_dict'] = self.optimizer.state_dict()
+        state_dict['fp32_groups'] = self.fp32_groups
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+        """
+        Loads a state_dict created by an earlier call to state_dict().
+        If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``,
+        whose parameters in turn came from ``model``, it is expected that the user
+        will call ``model.load_state_dict()`` before
+        ``fp16_optimizer_instance.load_state_dict()`` is called.
+        Example::
+            model = torch.nn.Linear(D_in, D_out).cuda().half()
+            optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
+            optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
+            ...
+            checkpoint = torch.load("saved.pth")
+            model.load_state_dict(checkpoint['model'])
+            optimizer.load_state_dict(checkpoint['optimizer'])
+        """
+        # I think it should actually be ok to reload the optimizer before the model.
+        self.dynamic_loss_scale = state_dict['dynamic_loss_scale']
+        self.cur_scale = state_dict['cur_scale']
+        self.cur_iter = state_dict['cur_iter']
+        if state_dict['dynamic_loss_scale']:
+            self.last_overflow_iter = state_dict['last_overflow_iter']
+            self.scale_factor = state_dict['scale_factor']
+            self.scale_window = state_dict['scale_window']
+        self.optimizer.load_state_dict(state_dict['optimizer_state_dict'])
+        # At this point, the optimizer's references to the model's fp32 parameters are up to date.
+        # The optimizer's hyperparameters and internal buffers are also up to date.
+        # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still
+        # out of date.  There are two options.
+        # 1:  Refresh the master params from the model's fp16 params.
+        # This requires less storage but incurs precision loss.
+        # 2:  Save and restore the fp32 master copies separately.
+        # We choose option 2.
+        #
+        # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device
+        # of their associated parameters, because it's possible those buffers might not exist yet in
+        # the current optimizer instance.  In our case, as long as the current FP16_Optimizer has been
+        # constructed in the same way as the one whose state_dict we are loading, the same master params
+        # are guaranteed to exist, so we can just copy_() from the saved master params.
+        for current, saved in zip(self.fp32_groups, state_dict['fp32_groups']):
+            for _current, _saved in zip(current, saved):
+                _current.data.copy_(_saved.data)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_adam.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_adam.py
new file mode 100644
index 0000000000..a823e7be6f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_adam.py
@@ -0,0 +1,206 @@
+import types
+import torch
+import importlib
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class FusedAdam(torch.optim.Optimizer):
+
+    """Implements Adam algorithm. Currently GPU-only.  Requires Apex to be installed via
+    ``python setup.py install --cuda_ext --cpp_ext``.
+
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square. (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False) NOT SUPPORTED in FusedAdam!
+        eps_inside_sqrt (boolean, optional): in the 'update parameters' step,
+            adds eps to the bias-corrected second moment estimate before
+            evaluating square root instead of adding it to the square root of
+            second moment estimate as in the original paper. (default: False)
+        use_mt (boolean, optional): use multi tensor apply for lower launch
+            latency. (default: False)
+
+    .. _Adam - A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params,
+                 lr=1e-3, bias_correction = True,
+                 betas=(0.9, 0.999), eps=1e-8, eps_inside_sqrt = False,
+                 weight_decay=0., max_grad_norm=0., amsgrad=False, use_mt=False,
+                 amp_scale_adjustment=1.0):
+        global fused_adam_cuda
+        fused_adam_cuda = importlib.import_module("fused_adam_cuda")
+
+        self._use_multi_tensor = False
+        if use_mt:
+            if not multi_tensor_applier.available:
+                print("Warning:  multi_tensor_applier is unavailable")
+            else:
+                self._use_multi_tensor = True
+                self._overflow_buf = torch.cuda.IntTensor([0])
+
+        self._amp_scale_adjustment = amp_scale_adjustment
+
+        if amsgrad:
+            raise RuntimeError('FusedAdam does not support the AMSGrad variant.')
+        defaults = dict(lr=lr, bias_correction=bias_correction,
+                        betas=betas, eps=eps, weight_decay=weight_decay,
+                        max_grad_norm=max_grad_norm)
+        super(FusedAdam, self).__init__(params, defaults)
+        self.eps_mode = 0 if  eps_inside_sqrt else 1
+
+    def step(self, closure=None, grads=None, output_params=None, scale=1., grad_norms=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+            grads (list of tensors, optional): weight gradient to use for the
+                optimizer update. If gradients have type torch.half, parameters
+                are expected to be in type torch.float. (default: None)
+            output params (list of tensors, optional): A reduced precision copy
+                of the updated weights written out in addition to the regular
+                updated weights. Have to be of same type as gradients. (default: None)
+            scale (float, optional): factor to divide gradient tensor values
+                by before applying to weights. (default: 1)
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        if hasattr(self, "_amp_stash"):
+            grads = self._amp_stash.grads
+            output_params = self._amp_stash.output_params
+            scale = self._amp_stash.scale*self._amp_scale_adjustment
+            grad_norms = self._amp_stash.grad_norms
+
+        if grads is None:
+            grads_group = [None]*len(self.param_groups)
+        # backward compatibility
+        # assuming a list/generator of parameter means single group
+        elif isinstance(grads, types.GeneratorType):
+            grads_group = [grads]
+        elif type(grads[0])!=list:
+            grads_group = [grads]
+        else:
+            grads_group = grads
+
+        if output_params is None:
+            output_params_group = [None]*len(self.param_groups)
+        elif isinstance(output_params, types.GeneratorType):
+            output_params_group = [output_params]
+        elif type(output_params[0])!=list:
+            output_params_group = [output_params]
+        else:
+            output_params_group = output_params
+
+        if grad_norms is None:
+            grad_norms = [None]*len(self.param_groups)
+
+        for group, grads_this_group, output_params_this_group, grad_norm in zip(self.param_groups, grads_group, output_params_group, grad_norms):
+            if grads_this_group is None:
+               grads_this_group = [None]*len(group['params'])
+            if output_params_this_group is None:
+               output_params_this_group = [None]*len(group['params'])
+
+            # compute combined scale factor for this group
+            combined_scale = scale
+            if group['max_grad_norm'] > 0:
+                # norm is in fact norm*scale
+                clip = ((grad_norm / scale) + 1e-6) / group['max_grad_norm']
+                if clip > 1:
+                    combined_scale = clip * scale
+
+            bias_correction = 1 if group['bias_correction'] else 0
+
+            if self._use_multi_tensor:
+                if output_params:
+                    tensorlists = [[],[],[],[],[]]
+                else:
+                    tensorlists = [[],[],[],[]]
+                tensordevice = None
+
+            for p, grad, output_param in zip(group['params'], grads_this_group, output_params_this_group):
+                #note: p.grad should not ever be set for correct operation of mixed precision optimizer that sometimes sends None gradients
+                if p.grad is None and grad is None:
+                    continue
+                if grad is None:
+                    grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('FusedAdam does not support sparse gradients, please consider SparseAdam instead')
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                out_p = torch.tensor([], dtype = torch.float) if output_param is None else output_param
+                if self._use_multi_tensor:
+                    pl = [p.data, exp_avg, exp_avg_sq, grad]
+                    if output_param is not None:
+                        pl.append(out_p)
+
+                    for tl, t in zip(tensorlists, pl):
+                        tl.append(t)
+
+                    if tensordevice is None:
+                        tensordevice = p.device
+                    elif tensordevice != p.device:
+                        raise RuntimeError('FusedAdam does not support use_mt with tensors on multiple device')
+
+                else:
+                    with torch.cuda.device(p.device):
+                        fused_adam_cuda.adam(p.data,
+                                             out_p,
+                                             exp_avg,
+                                             exp_avg_sq,
+                                             grad,
+                                             group['lr'],
+                                             beta1,
+                                             beta2,
+                                             group['eps'],
+                                             combined_scale,
+                                             state['step'],
+                                             self.eps_mode,
+                                             bias_correction,
+                                             group['weight_decay'])
+
+            if self._use_multi_tensor:
+                with torch.cuda.device(tensordevice):
+                    multi_tensor_applier(
+                        fused_adam_cuda.adam_mt,
+                        self._overflow_buf,
+                        tensorlists,
+                        group['lr'],
+                        beta1,
+                        beta2,
+                        group['eps'],
+                        combined_scale,
+                        state['step'],
+                        self.eps_mode,
+                        bias_correction,
+                        group['weight_decay'])
+
+        return loss
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_lamb.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_lamb.py
new file mode 100644
index 0000000000..81d8682282
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_lamb.py
@@ -0,0 +1,208 @@
+import torch
+import importlib
+import math
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class FusedLAMB(torch.optim.Optimizer):
+
+    """Implements LAMB algorithm.
+
+    Currently GPU-only.  Requires Apex to be installed via
+    ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" --global-option="--deprecated_fused_lamb" ./``.
+
+    This version of fused LAMB implements 2 fusions.
+
+      * Fusion of the LAMB update's elementwise operations
+      * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
+
+    :class:`apex.contrib.optimizers.FusedLAMB`'s usage is identical to any ordinary Pytorch optimizer::
+
+        opt = apex.contrib.optimizers.FusedLAMB(model.parameters(), lr = ....)
+        ...
+        opt.step()
+
+    :class:`apex.optimizers.FusedLAMB` may be used with or without Amp.  If you wish to use :class:`FusedLAMB` with Amp,
+    you may choose any ``opt_level``::
+
+        opt = apex.optimizers.FusedLAMB(model.parameters(), lr = ....)
+        model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
+        ...
+        opt.step()
+
+    In general, ``opt_level="O1"`` is recommended.
+
+    LAMB was proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its norm. (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            NOT SUPPORTED now! (default: False)
+        adam_w_mode (boolean, optional): Apply L2 regularization or weight decay
+            True for decoupled weight decay(also known as AdamW) (default: True)
+        grad_averaging (bool, optional): whether apply (1-beta2) to grad when
+            calculating running averages of gradient. (default: True)
+        set_grad_none (bool, optional): whether set grad to None when zero_grad()
+            method is called. (default: True)
+        max_grad_norm (float, optional): value used to clip global grad norm
+            (default: 1.0)
+
+    .. _Large Batch Optimization for Deep Learning - Training BERT in 76 minutes:
+        https://arxiv.org/abs/1904.00962
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params, lr=1e-3, bias_correction=True,
+                 betas=(0.9, 0.999), eps=1e-6, weight_decay=0.01,
+                 amsgrad=False, adam_w_mode=True,
+                 grad_averaging=True, set_grad_none=True,
+                 max_grad_norm=1.0):
+        if amsgrad:
+            raise RuntimeError('FusedLAMB does not support the AMSGrad variant.')
+        defaults = dict(lr=lr, bias_correction=bias_correction,
+                        betas=betas, eps=eps, weight_decay=weight_decay,
+                        grad_averaging=grad_averaging,
+                        max_grad_norm=max_grad_norm)
+        super(FusedLAMB, self).__init__(params, defaults)
+        if multi_tensor_applier.available:
+            import amp_C
+            self.multi_tensor_l2norm=amp_C.multi_tensor_l2norm
+            self._dummy_overflow_buf = torch.cuda.IntTensor([0])
+            fused_lamb_cuda = importlib.import_module("fused_lamb_cuda")
+            self.multi_tensor_lamb = fused_lamb_cuda.lamb
+        else:
+            raise RuntimeError('apex.contrib.optimizers.FusedLAMB requires cuda extensions')
+
+        self.adam_w_mode = 1 if adam_w_mode else 0
+        self.set_grad_none = set_grad_none
+
+    def zero_grad(self):
+        if self.set_grad_none:
+            for group in self.param_groups:
+                for p in group['params']:
+                    p.grad = None
+        else:
+            super(FusedLAMB, self).zero_grad()
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        # create separate grad lists for fp32 and fp16 params
+        g_all_32, g_all_16 = [], []
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                if p.dtype == torch.float32:
+                    g_all_32.append(p.grad.data)
+                elif p.dtype == torch.float16:
+                    g_all_16.append(p.grad.data)
+                else:
+                    raise RuntimeError('FusedLAMB only support fp16 and fp32.')
+
+        g_norm_32, g_norm_16 = 0.0, 0.0
+        # compute grad norm for two lists
+        if len(g_all_32) > 0:
+            g_norm_32 = multi_tensor_applier(self.multi_tensor_l2norm,
+                                             self._dummy_overflow_buf,
+                                             [g_all_32], False)[0].item()
+        if len(g_all_16) > 0:
+            g_norm_16 = multi_tensor_applier(self.multi_tensor_l2norm,
+                                             self._dummy_overflow_buf,
+                                             [g_all_16], False)[0].item()
+
+        # blend two grad norms to get global grad norm
+        global_grad_norm = math.sqrt(g_norm_32 * g_norm_32 + g_norm_16 * g_norm_16)
+        max_grad_norm = self.defaults['max_grad_norm']
+
+        for group in self.param_groups:
+            bias_correction = 1 if group['bias_correction'] else 0
+            beta1, beta2 = group['betas']
+            grad_averaging = 1 if group['grad_averaging'] else 0
+
+            # assume same step across group now to simplify things
+            # per parameter step can be easily support by making it tensor, or pass list into kernel
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            # create lists for multi-tensor apply
+            g_16, p_16, m_16, v_16 = [], [], [], []
+            g_32, p_32, m_32, v_32 = [], [], [], []
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                if p.grad.data.is_sparse:
+                    raise RuntimeError('FusedLAMB does not support sparse gradients, please consider SparseAdam instead')
+
+                state = self.state[p]
+                # State initialization
+                if len(state) == 0:
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+                if p.dtype == torch.float16:
+                    g_16.append(p.grad.data)
+                    p_16.append(p.data)
+                    m_16.append(state['exp_avg'])
+                    v_16.append(state['exp_avg_sq'])
+                elif p.dtype == torch.float32:
+                    g_32.append(p.grad.data)
+                    p_32.append(p.data)
+                    m_32.append(state['exp_avg'])
+                    v_32.append(state['exp_avg_sq'])
+                else:
+                    raise RuntimeError('FusedLAMB only support fp16 and fp32.')
+
+            if(len(g_16) > 0):
+                multi_tensor_applier(self.multi_tensor_lamb,
+                                     self._dummy_overflow_buf,
+                                     [g_16, p_16, m_16, v_16],
+                                     group['lr'],
+                                     beta1,
+                                     beta2,
+                                     group['eps'],
+                                     group['step'],
+                                     bias_correction,
+                                     group['weight_decay'],
+                                     grad_averaging,
+                                     self.adam_w_mode,
+                                     global_grad_norm,
+                                     max_grad_norm)
+            if(len(g_32) > 0):
+                multi_tensor_applier(self.multi_tensor_lamb,
+                                     self._dummy_overflow_buf,
+                                     [g_32, p_32, m_32, v_32],
+                                     group['lr'],
+                                     beta1,
+                                     beta2,
+                                     group['eps'],
+                                     group['step'],
+                                     bias_correction,
+                                     group['weight_decay'],
+                                     grad_averaging,
+                                     self.adam_w_mode,
+                                     global_grad_norm,
+                                     max_grad_norm)
+
+        return loss
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_sgd.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_sgd.py
new file mode 100644
index 0000000000..83587c6a63
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/optimizers/fused_sgd.py
@@ -0,0 +1,211 @@
+import types
+import torch
+from torch.optim.optimizer import Optimizer, required
+
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class FusedSGD(Optimizer):
+    r"""Implements stochastic gradient descent (optionally with momentum).
+
+    This version of fused SGD implements 2 fusions.
+      * Fusion of the SGD update's elementwise operations
+      * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
+
+    :class:`apex.contrib.optimizers.FusedSGD` should be used without AMP.
+   
+    :class:`apex.contrib.optimizers.FusedSGD` only works in the case where all parameters require grad. 
+
+    Nesterov momentum is based on the formula from
+    `On the importance of initialization and momentum in deep learning`__.
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float): learning rate
+        momentum (float, optional): momentum factor (default: 0)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        dampening (float, optional): dampening for momentum (default: 0)
+        nesterov (bool, optional): enables Nesterov momentum (default: False)
+
+    Example:
+        model = ...
+        model.half()
+        optimizer = apex.contrib.optimizers.FusedSGD(model.parameters())
+        # wrap with FP16_Optimizer
+        optimizer = FP16_Optimizer(optimizer, dynamic_loss_scale=True)
+        optimizer.zero_grad()
+	...
+        optimizer.backward(loss)
+        optmizer.step()
+
+    __ http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf
+
+    .. note::
+        The implementation of SGD with Momentum/Nesterov subtly differs from
+        Sutskever et. al. and implementations in some other frameworks.
+
+        Considering the specific case of Momentum, the update can be written as
+
+        .. math::
+                  v = \rho * v + g \\
+                  p = p - lr * v
+
+        where p, g, v and :math:`\rho` denote the parameters, gradient,
+        velocity, and momentum respectively.
+
+        This is in contrast to Sutskever et. al. and
+        other frameworks which employ an update of the form
+
+        .. math::
+             v = \rho * v + lr * g \\
+             p = p - v
+
+        The Nesterov version is analogously modified.
+    """
+
+    def __init__(self, params, lr=required, momentum=0, dampening=0,
+                 weight_decay=0, nesterov=False,
+                 wd_after_momentum=False,
+                 materialize_master_grads=True):
+        if lr is not required and lr < 0.0:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if momentum < 0.0:
+            raise ValueError("Invalid momentum value: {}".format(momentum))
+        if weight_decay < 0.0:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+
+        defaults = dict(lr=lr, momentum=momentum, dampening=dampening,
+                        weight_decay=weight_decay, nesterov=nesterov)
+        if nesterov and (momentum <= 0 or dampening != 0):
+            raise ValueError("Nesterov momentum requires a momentum and zero dampening")
+        super(FusedSGD, self).__init__(params, defaults)
+
+        self.wd_after_momentum = wd_after_momentum
+
+        if multi_tensor_applier.available:
+            import amp_C
+            # Skip buffer
+            self._dummy_overflow_buf = torch.cuda.IntTensor([0])
+            self.multi_tensor_sgd = amp_C.multi_tensor_sgd
+        else:
+            raise RuntimeError('apex.contrib.optimizers.FusedSGD requires cuda extensions')
+
+    def __setstate__(self, state):
+        super(FusedSGD, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('nesterov', False)
+
+    def get_momentums(self, params):
+        momentums = []
+        first_run = True
+        for p in params:
+            param_state = self.state[p]
+            # torch.optim.SGD initializes momentum in the main loop, we have
+            # to do it here, and track whether or not we've done so, so that
+            # momentum application can be skipped in the main kernel.
+            if 'momentum_buffer' not in param_state:
+                first_run = True
+                buf = param_state['momentum_buffer'] = torch.zeros_like(p.data)
+                momentums.append(buf)
+            else:
+                first_run = False
+                momentums.append(param_state['momentum_buffer'])
+        return momentums, first_run
+    
+    def step(self, closure=None, grads=None, output_params=None, scale=1., grad_norms=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+            grads (list of tensors, optional): weight gradient to use for the
+                optimizer update. If gradients have type torch.half, parameters
+                are expected to be in type torch.float. (default: None)
+            output_params (list of tensors, optional): A reduced precision copy
+                of the updated weights written out in addition to the regular
+                updated weights. Have to be of same type as gradients. (default: None)
+            scale (float, optional): factor to divide gradient tensor values
+                by before applying to weights. (default: 1)
+        """
+        if hasattr(self, "_amp_stash"):
+            raise RuntimeError('apex.contrib.optimizers.FusedSGD should not be used with AMP.')
+
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        if grads is None:
+            raise RuntimeError('apex.contrib.optimizers.FusedSGD must be wrapped \
+	                       with apex.contrib.optimizers.FP16_Optimizer \
+			       which provides grads.')
+        # backward compatibility
+        # assuming a list/generator of parameter means single group
+        elif isinstance(grads, types.GeneratorType):
+            grads_group = [grads]
+        elif type(grads[0])!=list:
+            grads_group = [grads]
+        else:
+            grads_group = grads
+
+        if output_params is None:
+            raise RuntimeError('apex.contrib.optimizers.FusedSGD must be wrapped \
+                               with apex.contrib.optimizers.FP16_Optimizer \
+                               which provides output_params.')
+        elif isinstance(output_params, types.GeneratorType):
+            output_params_group = [output_params]
+        elif type(output_params[0])!=list:
+            output_params_group = [output_params]
+        else:
+            output_params_group = output_params
+
+        for group, grads_this_group, output_params_this_group in zip(self.param_groups, 
+	                                                             grads_group, 
+                                                                     output_params_group):
+            if grads_this_group is None or output_params_this_group is None: 
+                raise RuntimeError('apex.contrib.optimizers.FusedSGD only works \
+                                    when all parameters require grad.')
+            
+            weight_decay = group['weight_decay']
+            momentum = group['momentum']
+            dampening = group['dampening']
+            nesterov = group['nesterov']
+            lr = group['lr']
+
+            first_runs = [True, True]
+            
+            # output_params_this_group: original weights (either fp16 or fp32)
+            # group['params']: master weights (fp32)
+
+            # grad_type, param_to_update_type, momentum_type, requires_fp16_model_copy
+            # fp32, fp32, fp32, No
+            fp32_grads = [g for (p, g) in zip(output_params_this_group, grads_this_group) if p.dtype == torch.float32]
+            fp32_params = [p2 for (p1, p2) in zip(output_params_this_group, group['params']) if p1.dtype == torch.float32]
+            fp32_momentums, first_runs[1] = self.get_momentums(fp32_params)
+            fp32_set = [fp32_grads, fp32_params, fp32_momentums]
+
+            # fp16, fp32, fp32, Yes
+            fp16_grads = [g for (p, g) in zip(output_params_this_group, grads_this_group) if p.dtype == torch.float16]
+            fp32_from_fp16_params = [p2 for (p1, p2) in zip(output_params_this_group, group['params']) if p1.dtype == torch.float16]
+            fp32_from_fp16_momentums, first_runs[0] = self.get_momentums(fp32_from_fp16_params)
+            fp16_params = [p1 for (p1, p2) in zip(output_params_this_group, group['params']) if p1.dtype == torch.float16]
+            fp16_set = [fp16_grads, fp32_from_fp16_params, fp32_from_fp16_momentums, fp16_params]
+
+            launch_sets = [fp16_set, fp32_set]
+
+            for launch_set, first_run in zip(launch_sets, first_runs):
+                assert len(launch_set[0]) == len(launch_set[1])
+                assert len(launch_set[0]) == len(launch_set[2])
+                if len(launch_set[0]) > 0:
+                    multi_tensor_applier(
+                        self.multi_tensor_sgd,
+                        self._dummy_overflow_buf,
+                        launch_set,
+                        weight_decay,
+                        momentum,
+                        dampening,
+                        lr,
+                        nesterov,
+                        first_run,
+                        self.wd_after_momentum,
+                        1.0/scale)
+
+        return loss
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/README.md
new file mode 100644
index 0000000000..9f540a6d50
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/README.md
@@ -0,0 +1,78 @@
+# Introduction to ASP
+
+This serves as a quick-start for ASP (Automatic SParsity), a tool that enables sparse training and inference for PyTorch models by adding 2 lines of Python.
+
+## Importing ASP
+```
+from apex.contrib.sparsity import ASP
+```
+
+## Initializing ASP
+
+Apart from the import statement, it is sufficient to add just the following line of code before the training phase to augment the model and the optimizer for sparse training/inference:
+```
+ASP.prune_trained_model(model, optimizer)
+```
+
+In the context of a typical PyTorch training loop, it might look like this:
+```
+ASP.prune_trained_model(model, optimizer)
+
+x, y = DataLoader(args)
+for epoch in range(epochs):
+    y_pred = model(x)
+    loss = loss_function(y_pred, y)
+    loss.backward()
+    optimizer.step()
+
+torch.save(...)
+```
+The `prune_trained_model` step calculates the sparse mask and applies it to the weights. This is done once, i.e., sparse locations in the weights matrix remain fixed after this step. 
+
+## Generate a Sparse Network
+
+The following approach serves as a guiding example on how to generate a pruned model that can use Sparse Tensor Cores in the NVIDIA Ampere Architecture. This approach generates a model for deployment, i.e. inference mode.
+
+```
+(1) Given a fully trained (dense) network, prune parameter values in a 2:4 sparse pattern.
+(2) Fine-tune  the  pruned  model  with  optimization  method  and  hyper-parameters (learning-rate, schedule, number of epochs, etc.) exactly as those used to obtain the trained model.
+(3) (If required) Quantize the model.
+```
+
+In code, below is a sketch on how to use ASP for this approach (steps 1 and 2 above).
+
+```
+
+model = define_model(..., pretrained=True) # define model architecture and load parameter tensors with trained values (by reading a trained checkpoint)
+criterion = ... # compare ground truth with model predition; use the same criterion as used to generate the dense trained model
+optimizer = ... # optimize model parameters; use the same optimizer as used to generate the dense trained model
+lr_scheduler = ... # learning rate scheduler; use the same schedule as used to generate the dense trained model
+
+from apex.contrib.sparsity import ASP     
+ASP.prune_trained_model(model, optimizer) #pruned a trained model
+
+x, y = DataLoader(args)
+for epoch in range(epochs): # train the pruned model for the same number of epochs as used to generate the dense trained model
+    y_pred = model(x)
+    loss = criterion(y_pred, y)
+    lr_scheduler.step()
+    loss.backward()
+    optimizer.step()
+
+torch.save(...) # saves the pruned checkpoint with sparsity masks 
+```
+
+## Non-Standard Usage
+
+If your goal is to easily perpare a network for accelerated inference, please follow the recipe above.  However, ASP can also be used to perform experiments in advanced techniques like training with sparsity from initialization. For example, in order to recompute the sparse mask in between training steps, use the following method:
+
+```
+ASP.compute_sparse_masks()
+```
+
+A more thorough example can be found in `./test/toy_problem.py`. 
+
+
+
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/__init__.py
new file mode 100644
index 0000000000..661fd4ae94
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/__init__.py
@@ -0,0 +1,2 @@
+from .sparse_masklib import create_mask
+from .asp import ASP
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/asp.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/asp.py
new file mode 100644
index 0000000000..17b1f34853
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/asp.py
@@ -0,0 +1,217 @@
+import types
+import torch
+from .sparse_masklib import create_mask
+
+torchvision_imported=True
+try:
+    import torchvision
+except ImportError:
+    print("[ASP][Warning] torchvision cannot be imported.")
+    torchvision_imported=False
+
+def eligible_modules(model, whitelist_layer_types, allowed_layer_names, disallowed_layer_names):
+    eligible_modules_list = []
+    for name, mod in model.named_modules():
+        if isinstance(mod, whitelist_layer_types) and name not in disallowed_layer_names:
+            if allowed_layer_names is not None and name not in allowed_layer_names:
+                continue
+            eligible_modules_list.append((name, mod))
+    return eligible_modules_list
+
+class ASP:
+    __model = None
+    __verbosity = 0
+    __optimizer = None
+    __sparse_parameters = []
+    __calculate_mask = None
+
+    @classmethod
+    def init_model_for_pruning(cls, model, mask_calculator="m4n2_1d",
+             verbosity=3,
+             whitelist=[torch.nn.Linear, torch.nn.Conv1d, torch.nn.Conv2d, torch.nn.Conv3d], 
+             allowed_layer_names=None, disallowed_layer_names=[],
+             allow_recompute_mask=False, custom_layer_dict={}):
+        """Call this method to modify your model to take advantage of sparse matrix multiplication.
+        Note that this call alone only augments the model with additional buffers needed for sparse MMA,
+        it does not enable use of sparse MMA. 
+
+        If you are starting with a fresh model:
+
+        model = ...
+        ASP.init_model_for_pruning(model, mask_calculator, ...)
+        if (training) ASP.init_optimizer_for_pruning(optimizer)
+        ASP.compute_sparse_masks() // sparsity is off by default, call when youy want to enable it.
+
+        If you are starting from a checkpoint:
+
+        model = ...
+        ASP.init_model_for_pruning(model, mask_calculator, ...)
+        torch.load(...)
+        if (training) ASP.init_optimizer_for_pruning(optimizer)
+
+        Arguments:
+          model                    The model
+          mask_calculator          Either callable that computes mask given a tensor OR pattern string for sparse mask lib.
+          verbosity                Integer controling verbosity level.
+                                   0 -> Only errors.
+                                   1 -> Errors and warnings.
+                                   2 -> Errors, warnings and info.
+                                   3 -> Errors, warnings, info and debug.
+          whitelist                Module types approved for sparsity.
+          allowed_layer_names      If not None, only layer names that appear in this list are considered for sparsity.
+          disallowed_layer_names   If not [], only layer names that do not appear in this list are considered for sparsity.
+          allow_recompute_mask     If True, stores pruned values so that dense weights can be restored.
+                                   Pruned weights are stored in CPU memory, hence this option does not increase GPU memory usage.
+          custom_layer_dict        Dictionary of additional layer paremeters to sparsify. e.g. {CustomLinear: ['weight']}
+          
+          [Future] Support for allow_recompute_mask can be removed, it is not part of sparse inference recipe -- AKM. 
+        """
+        assert (cls.__model is None), "ASP has been initialized already."
+        cls.__model = model
+        cls.__verbosity = verbosity
+
+        if isinstance(mask_calculator, str):
+            def create_mask_from_pattern(param):
+                return create_mask(param, mask_calculator).bool()
+            cls.__calculate_mask = create_mask_from_pattern
+        else:
+            cls.__calculate_mask = mask_calculator #user defined function
+
+        # function to extract variables that will be sparsified. 
+        # idea is that you will add one of these functions for each module type that can be sparsified.
+        if torchvision_imported:
+            print("[ASP] torchvision is imported, can work with the MaskRCNN/KeypointRCNN from torchvision.")
+            sparse_parameter_list = {torch.nn.Linear: ['weight'], torch.nn.Conv1d: ['weight'], torch.nn.Conv2d: ['weight'], torch.nn.Conv3d: ['weight'], torchvision.ops.misc.Conv2d: ['weight']}
+        else:
+            sparse_parameter_list = {torch.nn.Linear: ['weight'], torch.nn.Conv1d: ['weight'], torch.nn.Conv2d: ['weight'], torch.nn.Conv3d: ['weight']}
+        if custom_layer_dict: # Update default list to include user supplied custom (layer type : parameter tensor), make sure this tensor type is something ASP knows how to prune
+            sparse_parameter_list.update(custom_layer_dict)
+            whitelist += list(custom_layer_dict.keys())
+
+        for module_type in whitelist:
+            assert (module_type in sparse_parameter_list), "Module %s :: Don't know how to sparsify module." % module.dtype()
+
+        # find all sparse modules, extract sparse parameters and decorate
+        def add_sparse_attributes(module_name, module):
+            sparse_parameters = sparse_parameter_list[type(module)]
+            for p_name, p in module.named_parameters():
+                if p_name in sparse_parameters and p.requires_grad:
+                    # check for NVIDIA's TC compatibility: we check along the horizontal direction
+                    if p.dtype == torch.float32 and ((p.size()[0] % 8) != 0 or (p.size()[1] % 16) != 0): #User defines FP32 and APEX internally uses FP16 math
+                        print("[ASP] Auto skipping pruning %s::%s of size=%s and type=%s for sparsity" % (module_name, p_name, str(p.size()), str(p.dtype)))
+                        continue
+                    if p.dtype == torch.float16 and ((p.size()[0] % 8) != 0 or (p.size()[1] % 16) != 0): #For Conv2d dim= K x CRS; we prune along C
+                        print("[ASP] Auto skipping pruning %s::%s of size=%s and type=%s for sparsity" % (module_name, p_name, str(p.size()), str(p.dtype)))
+                        continue
+                    
+                    if cls.__verbosity >= 3:
+                        print("[ASP] Sparsifying %s::%s of size=%s and type=%s for sparsity" % (module_name, p_name, str(p.size()), str(p.dtype)))
+                    
+                    mask = torch.ones_like(p).bool()
+                    buffname = p_name.split(".")[-1] # buffer names cannot contain "."
+                    module.register_buffer('__%s_mma_mask' % buffname, mask)
+                    if allow_recompute_mask:
+                        pruned = torch.zeros_like(p).cpu()
+                        module.register_buffer('__%s_mma_pruned_p' % buffname, pruned)
+                    else:
+                        pruned = None
+                    cls.__sparse_parameters.append((module_name, module, p_name, p, mask, pruned))
+                else:
+                    if cls.__verbosity >= 3:
+                        print("[ASP] Not sparsifying %s::%s of size=%s and type=%s" % (module_name, p_name, str(p.size()), str(p.dtype)))
+
+        for name, sparse_module in eligible_modules(model, tuple(whitelist), allowed_layer_names, disallowed_layer_names):
+            add_sparse_attributes(name, sparse_module)
+
+    @classmethod
+    def init_optimizer_for_pruning(cls, optimizer):
+        """Call this method to monkey patch optimizer step function so that masks can be applied to
+        gradients and weights during training.
+        You must call init_model_for_pruning(...) before calling init_optimizer_for_pruning(...)
+        """
+        assert (cls.__optimizer is None), "ASP has initialized optimizer already."
+        assert (cls.__calculate_mask is not None), "Called ASP.init_optimizer_for_pruning before ASP.init_model_for_pruning."
+
+        # store pointer to original optimizer step method
+        cls.__optimizer = optimizer
+        cls.__optimizer.__step = optimizer.step
+
+        def __step(opt_self, *args, **kwargs):
+            # prune gradients before step method
+            with torch.no_grad():
+                for module_name, module, p_name, p, mask, pruned in cls.__sparse_parameters:
+                    if p.grad is not None: #thx pjudd
+                        p.grad.mul_(mask)
+            # call original optimizer step method
+            rval = opt_self.__step(*args, **kwargs)
+            # prune parameters after step method
+            with torch.no_grad():
+                for module_name, module, p_name, p, mask, pruned in cls.__sparse_parameters:
+                    p.mul_(mask)
+            return rval
+        cls.__optimizer.step = types.MethodType(__step, cls.__optimizer)
+
+    @classmethod
+    def compute_sparse_masks(cls):
+        """Call this method to enable sparsity.
+        If init(...) was called with allow_recompute_mask=False AND sparsity is disabled, pruned field can be None.
+        """
+        with torch.no_grad():
+            for module_name, module, p_name, p, mask, pruned in cls.__sparse_parameters:
+                if mask.sum() < mask.numel(): # when recalculating masks
+                    # restore dense parameter if allow_recompute_mask is enabled
+                    assert (pruned is not None), "Unable to restore dense parameter because allow_recompute_mask == False"
+                    p.add_(pruned.cuda())
+
+                mask.set_(cls.__calculate_mask(p))
+
+                if pruned is not None: # stow away pruned weights to cpu
+                    pruned.set_((p * (~mask)).cpu())
+
+                p.mul_(mask) # in-place multiplication, so pruned weights are 0-values, hence checkpoint will have 0s for pruned weights
+                if cls.__verbosity >= 2:
+                    print("[ASP] Enabled %.2f%% sparsity for %s::%s of size=%s and type=%s" % (100.0*mask.sum()/mask.numel(), module_name, p_name, str(p.size()), str(p.dtype)))
+
+    @classmethod
+    def restore_pruned_weights(cls):
+        """Call this method to disable sparsity and restore all weights.
+        This will only work if init(...) was called with allow_recompute=True.
+        """
+        with torch.no_grad():
+            for module_name, module, p_name, p, mask, pruned in cls.__sparse_parameters:
+                if mask.sum() < mask.numel():
+                    assert (pruned is not None), "Unable to restore dense parameter because allow_recompute_mask == False"
+                    p.add_(pruned.cuda())
+                    mask.fill_(1)
+                    pruned.zero_()
+                    if cls.__verbosity >= 2:
+                        print("[ASP] Disabled sparsity for %s::%s (dense weights restored)" % (module_name, p_name))
+
+    @classmethod
+    def is_sparsity_enabled(cls):
+        """Call this method to determine if sparsity is enabled in the model.
+        The typical use case is right after checkpoint has been loaded.
+        """
+        total,sp100,sp50 = 0,0,0
+        for module_name, module, p_name, p, mask, pruned in cls.__sparse_parameters:
+            total += 1
+            mask_sum = mask.sum()
+            mask_numel = mask.numel()
+            if mask_sum == mask_numel:
+                sp100 += 1
+            elif mask_sum*2 == mask_numel:
+                sp50 += 1
+
+        assert (total == sp100 or total == sp50), "Inconsistent model sparsity"
+        if total == sp100:
+            return False
+        elif total == sp50:
+            return True
+    
+    @classmethod
+    def prune_trained_model(cls, model, optimizer):
+        # add mask buffers to model (init_model_for_pruning), augment optimizer (init_optimizer_for_pruning) and compute masks (compute_sparse_masks)
+        cls.init_model_for_pruning(model, mask_calculator="m4n2_1d", verbosity=2, whitelist=[torch.nn.Linear, torch.nn.Conv2d], allow_recompute_mask=False)
+        cls.init_optimizer_for_pruning(optimizer)
+        cls.compute_sparse_masks()
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/sparse_masklib.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/sparse_masklib.py
new file mode 100644
index 0000000000..ed42d0456d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/sparse_masklib.py
@@ -0,0 +1,184 @@
+import sys
+import torch
+import numpy as np
+import collections
+from itertools import permutations
+
+
+""" compute density (helper fn to compute % NNZs in a tensor) """
+def fill(x):
+    return float(x.nonzero().size(0))/torch.numel(x)
+
+""" reshape matrix into m-dimensional vectors: (h,w) -> (hw/m, m) """
+def reshape_1d(matrix, m):
+    # If not a nice multiple of m, fill with zeroes.
+    if matrix.shape[1] % m > 0:
+        mat = torch.cuda.FloatTensor(matrix.shape[0], matrix.shape[1] + (m-matrix.shape[1]%m)).fill_(0)
+        mat[:, :matrix.shape[1]] = matrix
+        shape = mat.shape
+        return mat.view(-1,m),shape
+    else:
+        return matrix.view(-1,m), matrix.shape
+
+""" return all possible m:n patterns in a 1d vector """
+valid_m4n2_1d_patterns = None
+def compute_valid_1d_patterns(m,n):
+    # Early exit if patterns was already created.
+    global valid_m4n2_1d_patterns
+
+    if m==4  and n==2 and valid_m4n2_1d_patterns  is not None: return valid_m4n2_1d_patterns
+    patterns = torch.zeros(m)
+    patterns[:n] = 1
+    valid_patterns = torch.Tensor(list(set(permutations(patterns.tolist()))))
+    if m == 4  and n == 2: valid_m4n2_1d_patterns  = valid_patterns       
+    return valid_patterns
+
+""" m:n 1d structured best """
+def mn_1d_best(matrix, m, n):
+    # Find all possible patterns.
+    patterns = compute_valid_1d_patterns(m,n).cuda()
+
+    # Find the best m:n pattern (sum of non-masked weights).
+    mask = torch.cuda.IntTensor(matrix.shape).fill_(1).view(-1,m)
+    mat,shape = reshape_1d(matrix,m)
+    pmax = torch.argmax(torch.matmul(mat.abs(),patterns.t()), dim=1)
+    mask[:] = patterns[pmax[:]]
+    mask = mask.view(matrix.shape)
+    return mask
+
+def m4n2_1d(mat, density):
+    return mn_1d_best(mat, 4, 2)
+
+"""
+  Below 2d-masking related code is targeted more for training (from scratch).
+  2d-pruning of a weight tensor is done to accelerate DGRAD step during backprop
+  phase of training algorithm. Acceleration comes from using SpMMA instructions in
+  Tensor Cores of NVIDIA Ampere GPU Architecture 
+  (note: this code does not do the acceleration, GPU kernels are required for this).
+  1d pruning of weight tensor helps speed up FPROP step by pruning in 2:4 pattern
+  along the horizontal (logical) direction.
+  During DGRAD step, weight tensor is transposed. 2d pruning functions below, mask
+  weight tensor such that their transposed versions are also 2:4 sparse along the
+  horizontal (logical) direction. Thus, with 2d pruning, weight tensors are 
+  2:4 sparse along row and column directions.
+ """
+
+""" m:n 2d structured pruning: greedy method to select mask """
+def mn_2d_greedy(matrix, m, n):
+    # Convert to numpy
+    mat = matrix.cpu().detach().numpy()
+    mask = np.ones(mat.shape, dtype=int)
+
+    rowCount = int(mat.shape[0]/m) * m
+    colCount = int(mat.shape[1]/m) * m
+    for rowStartIdx in range(0, rowCount, m):
+        rowEndIdx = rowStartIdx + m
+        for colStartIdx in range(0, colCount, m):
+            colEndIdx = colStartIdx + m
+            matrixSub = np.absolute(np.squeeze(mat[rowStartIdx:rowEndIdx, colStartIdx:colEndIdx]))
+            maskSub = np.squeeze(mask[rowStartIdx:rowEndIdx, colStartIdx:colEndIdx])
+            maskSub.fill(0.0)
+            matrixVecView = matrixSub.reshape(-1)
+            maskVecView   = maskSub.reshape(-1)
+            linearIdx = np.argsort(matrixVecView)
+            matrixIdx = [(int(x/m), x % m) for x in linearIdx]
+            rowCounter = collections.Counter()
+            colCounter = collections.Counter()
+            for currIdx in range(len(linearIdx) - 1, -1, -1):
+                currMatrixEntry = matrixIdx[currIdx]
+                if (rowCounter[currMatrixEntry[0]] == n) or (colCounter[currMatrixEntry[1]] == n):
+                    continue
+                #end if
+                maskSub[currMatrixEntry[0], currMatrixEntry[1]] = 1.0
+                rowCounter[currMatrixEntry[0]] += 1
+                colCounter[currMatrixEntry[1]] += 1
+
+    return torch.tensor(mask.cuda())
+
+def m4n2_2d_greedy(mat, density):
+    return mn_2d_greedy(mat, 4, 2)
+
+""" return all possible m:n patterns in a mxn block. """
+valid_m4n2_2d_patterns = None
+def compute_valid_2d_patterns(m,n):
+    # Early exit if patterns was already created.
+    global valid_m4n2_2d_patterns
+    if valid_m4n2_2d_patterns is not None: return valid_m4n2_2d_patterns
+
+    patterns = torch.zeros(m)
+    patterns[:n] = 1
+    patterns = list(set(permutations(patterns.tolist())))
+    patterns = patterns + patterns
+    patterns = torch.Tensor(list(set(permutations(patterns,m))))
+
+    valid = ((patterns.sum(dim=1) <= n).sum(dim=1) == m).nonzero().view(-1)
+    valid_patterns = torch.Tensor(valid.shape[0],m,m)
+    valid_patterns[:] = patterns[valid[:]]
+
+    if m == 4  and n == 2: valid_m4n2_2d_patterns  = valid_patterns
+    return valid_patterns
+
+""" m:n 2d structured pruning: exhaustive method to select best mask """
+def mn_2d_best(matrix, m, n):
+    # Find all possible patterns.
+    patterns = compute_valid_2d_patterns(m,n).cuda()
+
+    # Find the best m:n pattern (sum of non-masked weights).
+    mask = torch.cuda.IntTensor(matrix.shape).fill_(1)
+    mat = reshape_2d(matrix,m,m).abs()
+    pmax = torch.argmax(torch.matmul(mat,patterns.view(patterns.shape[0],m*m).t()), dim=2)
+
+    # Copy best m:n patterns into mask.
+    mat = mat.view(mat.shape[0]*mat.shape[1],-1)
+    pmax = pmax.view(pmax.shape[0]*pmax.shape[1]).unsqueeze(1).expand(-1,mat.shape[1])
+    patterns = patterns.view(patterns.shape[0],patterns.shape[1]*patterns.shape[2])
+    mat = torch.gather(patterns,0,pmax)
+    mat = reshape_2d_inv(mat.view(matrix.shape[0]//m,matrix.shape[1]//m,m,m))
+    mask.copy_(mat.type(mask.type()))
+    return mask
+
+def m4n2_2d_best(mat, density):
+    return mn_2d_best(mat, 4, 2)
+
+
+""" returns a sparse mask """
+def create_mask(tensor, pattern="m4n2_1d", density=0.5):
+    # Reshape tensor and mask.
+    shape = tensor.shape
+    ttype = tensor.type()
+    t = tensor.float().contiguous()
+
+    # 1d-tensor
+    if len(shape) == 1:
+        t = t.view(1, shape[0])
+        func = getattr(sys.modules[__name__], pattern, None)
+        mask = func(t, density)
+        return mask.view(shape).type(ttype)
+    # 2d-tensor (in, out)
+    elif len(shape) == 2:
+        t = t.view(shape[0], shape[1])
+        func = getattr(sys.modules[__name__], pattern, None)
+        mask = func(t, density)
+        return mask.view(shape).type(ttype)
+    # 3d-tensor (batch, in, out)
+    elif len(shape) == 3:
+        t = t.view(shape[0]*shape[1], shape[2])
+        func = getattr(sys.modules[__name__], pattern, None)
+        mask = func(t, density)
+        return mask.view(shape).type(ttype)
+    # 4d-tensor (in, out, h, w)
+    elif len(shape) == 4:
+        """
+        # transformers (bmm)
+        t = t.view(shape[0]*shape[1]*shape[2], shape[3])
+        func = getattr(sys.modules[__name__], pattern, None)
+        mask = func(t, density)
+        return mask.view(shape).type(ttype)
+        """
+        # convs
+        t = t.permute(2,3,0,1).contiguous().view(shape[2]*shape[3]*shape[0], shape[1])
+        func = getattr(sys.modules[__name__], pattern, None)
+        mask = func(t, density)
+        mask = mask.view(shape[2], shape[3], shape[0], shape[1]).permute(2,3,0,1).contiguous()      
+        return mask.view(shape).type(ttype)
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_part1.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_part1.py
new file mode 100644
index 0000000000..67350fea2e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_part1.py
@@ -0,0 +1,94 @@
+from collections import OrderedDict
+
+import torch
+from apex.optimizers import FusedAdam
+from apex.contrib.sparsity import ASP
+
+def build_model(args):
+    od = OrderedDict()
+    for i in range(args.num_layers):
+        if i == 0:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.input_features, args.hidden_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.hidden_features])
+        elif i == args.num_layers-1:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.hidden_features, args.output_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.output_features])
+        else:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.hidden_features, args.hidden_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.hidden_features])
+    return torch.nn.Sequential(od)
+
+def train_step(args, model, optimizer, input_batch, target_batch, step):
+    predicted_target = model(input_batch)
+    loss = ((predicted_target-target_batch)**2).sum()
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+    step = step + 1
+    #print("Step %d :: loss=%e" % (step, loss.item()))
+    return step
+
+def train_loop(args, model, optimizer, step, num_steps):
+    for i in range(num_steps):
+        input_batch = torch.randn([args.batch_size, args.input_features]).cuda()
+        target_batch = torch.randn([args.batch_size, args.output_features]).cuda()
+        step = train_step(args, model, optimizer, input_batch, target_batch, step)
+    return step
+
+def main(args):
+    #
+    # PART1
+    #
+
+    torch.manual_seed(args.seed)
+
+    model = build_model(args).cuda()
+    one_ll = next(model.children()).weight
+    optimizer = FusedAdam(model.parameters())
+    ASP.init_model_for_pruning(model, args.pattern, verbosity=args.verbosity, whitelist=args.whitelist, allow_recompute_mask=args.allow_recompute_mask)
+    ASP.init_optimizer_for_pruning(optimizer)
+
+    step = 0
+
+    # train for a few steps with dense weights
+    print("DENSE :: ",one_ll)
+    step = train_loop(args, model, optimizer, step, args.num_dense_steps)
+
+    # simulate sparsity by inserting zeros into existing dense weights
+    ASP.enable_sparsity()
+
+    # train for a few steps with sparse weights
+    print("SPARSE :: ",one_ll)
+    step = train_loop(args, model, optimizer, step, args.num_sparse_steps)
+
+    torch.save({
+            'step': step,
+            'verbosity': args.verbosity,
+            'seed2': args.seed2,
+            'pattern': args.pattern,
+            'whitelist': args.whitelist,
+            'allow_recompute_mask': args.allow_recompute_mask,
+            'model_state_dict': model.state_dict(),
+            'optimizer_state_dict': optimizer.state_dict(),
+            }, args.checkpoint_path)
+
+if __name__ == '__main__':
+    class Args:
+        verbosity=3
+        seed = 4873
+        seed2 = 99875
+        pattern = "m4n2_2d_best"
+        whitelist = [torch.nn.Linear]
+        allow_recompute_mask = True
+        batch_size = 32
+        input_features = 8
+        output_features = 8
+        hidden_features = 32
+        num_layers = 4
+        num_dense_steps = 2000
+        num_sparse_steps = 3000
+        num_sparse_steps_2 = 1000
+        checkpoint_path = "part1.chkp"
+    args = Args()
+
+    main(args)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_part2.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_part2.py
new file mode 100644
index 0000000000..2ac24526f0
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_part2.py
@@ -0,0 +1,79 @@
+from collections import OrderedDict
+
+import torch
+from apex.optimizers import FusedAdam
+from apex.contrib.sparsity import ASP
+
+def build_model(args):
+    od = OrderedDict()
+    for i in range(args.num_layers):
+        if i == 0:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.input_features, args.hidden_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.hidden_features])
+        elif i == args.num_layers-1:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.hidden_features, args.output_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.output_features])
+        else:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.hidden_features, args.hidden_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.hidden_features])
+    return torch.nn.Sequential(od)
+
+def train_step(args, model, optimizer, input_batch, target_batch, step):
+    predicted_target = model(input_batch)
+    loss = ((predicted_target-target_batch)**2).sum()
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+    step = step + 1
+    #print("Step %d :: loss=%e" % (step, loss.item()))
+    return step
+
+def train_loop(args, model, optimizer, step, num_steps):
+    for i in range(num_steps):
+        input_batch = torch.randn([args.batch_size, args.input_features]).cuda()
+        target_batch = torch.randn([args.batch_size, args.output_features]).cuda()
+        step = train_step(args, model, optimizer, input_batch, target_batch, step)
+    return step
+
+def main(step, args, model_state_dict, optimizer_state_dict):
+    #
+    # PART2
+    #
+
+    model = build_model(args).cuda()
+    one_ll = next(model.children()).weight
+    optimizer = FusedAdam(model.parameters())
+    ASP.init_model_for_pruning(model, args.pattern, verbosity=args.verbosity, whitelist=args.whitelist, allow_recompute_mask=args.allow_recompute_mask)
+    ASP.init_optimizer_for_pruning(optimizer)
+
+    torch.manual_seed(args.seed2)
+    model.load_state_dict(model_state_dict)
+    optimizer.load_state_dict(optimizer_state_dict)
+
+    print("Model sparsity is %s" % ("enabled" if ASP.sparsity_is_enabled() else "disabled"))
+
+    # train for a few steps with sparse weights
+    print("SPARSE :: ",one_ll)
+    step = train_loop(args, model, optimizer, step, args.num_sparse_steps_2)
+
+if __name__ == '__main__':
+    checkpoint = torch.load("part1.chkp")
+    class Args:
+        verbosity = checkpoint['verbosity']
+        seed = 4873
+        seed2 = checkpoint['seed2']
+        pattern = checkpoint['pattern']
+        whitelist = checkpoint['whitelist']
+        allow_recompute_mask = checkpoint['allow_recompute_mask']
+        batch_size = 32
+        input_features = 8
+        output_features = 8
+        hidden_features = 32
+        num_layers = 4
+        num_dense_steps = 2000
+        num_sparse_steps = 3000
+        num_sparse_steps_2 = 1000
+        checkpoint_path = "part1.chkp"
+    args = Args()
+
+    main(checkpoint['step'], args, checkpoint['model_state_dict'], checkpoint['optimizer_state_dict'])
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_reference.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_reference.py
new file mode 100644
index 0000000000..96e2b8cc00
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/checkpointing_test_reference.py
@@ -0,0 +1,96 @@
+from collections import OrderedDict
+
+import torch
+from apex.optimizers import FusedAdam
+from apex.contrib.sparsity import ASP
+
+#
+# Reference run for checkpointing test (part1 + part2)
+#
+
+def build_model(args):
+    od = OrderedDict()
+    for i in range(args.num_layers):
+        if i == 0:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.input_features, args.hidden_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.hidden_features])
+        elif i == args.num_layers-1:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.hidden_features, args.output_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.output_features])
+        else:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.hidden_features, args.hidden_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.hidden_features])
+    return torch.nn.Sequential(od)
+
+def train_step(args, model, optimizer, input_batch, target_batch, step):
+    predicted_target = model(input_batch)
+    loss = ((predicted_target-target_batch)**2).sum()
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+    step = step + 1
+    #print("Step %d :: loss=%e" % (step, loss.item()))
+    return step
+
+def train_loop(args, model, optimizer, step, num_steps):
+    for i in range(num_steps):
+        input_batch = torch.randn([args.batch_size, args.input_features]).cuda()
+        target_batch = torch.randn([args.batch_size, args.output_features]).cuda()
+        step = train_step(args, model, optimizer, input_batch, target_batch, step)
+    return step
+
+def main(args):
+    #
+    # PART1
+    #
+
+    torch.manual_seed(args.seed)
+
+    model = build_model(args).cuda()
+    one_ll = next(model.children()).weight
+    optimizer = FusedAdam(model.parameters())
+    ASP.init_model_for_pruning(model, args.pattern, whitelist=args.whitelist, allow_recompute_mask=args.allow_recompute_mask)
+    ASP.init_optimizer_for_pruning(optimizer)
+
+    step = 0
+
+    # train for a few steps with dense weights
+    print("DENSE :: ",one_ll)
+    step = train_loop(args, model, optimizer, step, args.num_dense_steps)
+
+    # simulate sparsity by inserting zeros into existing dense weights
+    ASP.enable_sparsity()
+
+    # train for a few steps with sparse weights
+    print("SPARSE :: ",one_ll)
+    step = train_loop(args, model, optimizer, step, args.num_sparse_steps)
+
+    #
+    # PART 2
+    #
+
+    torch.manual_seed(args.seed2)
+
+    # train for a few steps with sparse weights
+    print("SPARSE :: ",one_ll)
+    step = train_loop(args, model, optimizer, step, args.num_sparse_steps_2)
+
+if __name__ == '__main__':
+    class Args:
+        seed = 4873
+        seed2 = 99875
+        pattern = "m4n2_2d_best"
+        whitelist = [torch.nn.Linear]
+        allow_recompute_mask = True
+        batch_size = 32
+        input_features = 8
+        output_features = 8
+        hidden_features = 32
+        num_layers = 4
+        num_dense_steps = 2000
+        num_sparse_steps = 3000
+        num_sparse_steps_2 = 1000
+        checkpoint_path = "part1.chkp"
+    args = Args()
+
+    main(args)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/toy_problem.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/toy_problem.py
new file mode 100644
index 0000000000..2145323ddc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/sparsity/test/toy_problem.py
@@ -0,0 +1,87 @@
+from collections import OrderedDict
+
+import torch
+from apex.optimizers import FusedAdam
+from apex.contrib.sparsity import ASP
+
+def build_model(args):
+    od = OrderedDict()
+    for i in range(args.num_layers):
+        if i == 0:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.input_features, args.hidden_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.hidden_features])
+        elif i == args.num_layers-1:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.hidden_features, args.output_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.output_features])
+        else:
+            od['linear_layer_%d' % (i+1)] = torch.nn.Linear(args.hidden_features, args.hidden_features)
+            od['layer_norm_%d' % (i+1)] = torch.nn.LayerNorm([args.batch_size, args.hidden_features])
+    return torch.nn.Sequential(od)
+
+def train_step(args, model, optimizer, input_batch, target_batch, step):
+    predicted_target = model(input_batch)
+    loss = ((predicted_target-target_batch)**2).sum()
+    loss.backward()
+    optimizer.step()
+    optimizer.zero_grad()
+    step = step + 1
+    #print("Step %d :: loss=%e" % (step, loss.item()))
+    return step
+
+def train_loop(args, model, optimizer, step, num_steps):
+    for i in range(num_steps):
+        input_batch = torch.randn([args.batch_size, args.input_features]).cuda()
+        target_batch = torch.randn([args.batch_size, args.output_features]).cuda()
+        step = train_step(args, model, optimizer, input_batch, target_batch, step)
+    return step
+
+def main(args):
+    model = build_model(args).cuda()
+    one_ll = next(model.children()).weight
+    optimizer = FusedAdam(model.parameters())
+    # only prune linear layers, even though we also support conv1d, conv2d and conv3d
+    ASP.init_model_for_pruning(model, "m4n2_1d", whitelist=[torch.nn.Linear], allow_recompute_mask=True)
+    ASP.init_optimizer_for_pruning(optimizer)
+
+    step = 0
+
+    # train for a few steps with dense weights
+    print("DENSE :: ",one_ll)
+    step = train_loop(args, model, optimizer, step, args.num_dense_steps)
+
+    # simulate sparsity by inserting zeros into existing dense weights
+    ASP.compute_sparse_masks()
+
+    # train for a few steps with sparse weights
+    print("SPARSE :: ",one_ll)
+    step = train_loop(args, model, optimizer, step, args.num_sparse_steps)
+
+    # recompute sparse masks
+    ASP.compute_sparse_masks()
+
+    # train for a few steps with sparse weights
+    print("SPARSE :: ",one_ll)
+    step = train_loop(args, model, optimizer, step, args.num_sparse_steps_2)
+
+    # turn off sparsity
+    print("SPARSE :: ",one_ll)
+    ASP.restore_pruned_weights()
+
+    # train for a few steps with dense weights
+    print("DENSE :: ",one_ll)
+    step = train_loop(args, model, optimizer, step, args.num_dense_steps_2)
+
+if __name__ == '__main__':
+    class Args:
+        batch_size = 32
+        input_features = 16
+        output_features = 8
+        hidden_features = 40
+        num_layers = 4
+        num_dense_steps = 2000
+        num_sparse_steps = 3000
+        num_sparse_steps_2 = 1000
+        num_dense_steps_2 = 1500
+    args = Args()
+
+    main(args)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/fmha/test_fmha.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/fmha/test_fmha.py
new file mode 100644
index 0000000000..f930fe2857
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/fmha/test_fmha.py
@@ -0,0 +1,121 @@
+###############################################################################
+# Copyright (c) 2011-2021, NVIDIA CORPORATION.  All rights reserved.
+# 
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions are met:
+#     * Redistributions of source code must retain the above copyright
+#       notice, this list of conditions and the following disclaimer.
+#     * Redistributions in binary form must reproduce the above copyright
+#       notice, this list of conditions and the following disclaimer in the
+#       documentation and/or other materials provided with the distribution.
+#     * Neither the name of the NVIDIA CORPORATION nor the
+#       names of its contributors may be used to endorse or promote products
+#       derived from this software without specific prior written permission.
+# 
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+# ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+# WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+# DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+# DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+# (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+# LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+# ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+# SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+###############################################################################
+
+
+import sys
+import torch
+import numpy as np
+import unittest
+import math
+
+import fmhalib as mha
+
+def py_mha(qkv, amask, b, s, h, d):
+    qkv = qkv.view(b, s, h, 3, d)
+    q = qkv[:, :, :, 0, :].permute(0,2,1,3)
+    k = qkv[:, :, :, 1, :].permute(0,2,1,3)
+    v = qkv[:, :, :, 2, :].permute(0,2,1,3)
+    p = torch.matmul(q.float(), k.permute(0,1,3,2).float())
+    p_masked = p / math.sqrt(d) + (1.0 - amask) * -10000.0
+    s = torch.softmax(p_masked, -1).to(qkv.dtype)
+    ctx = torch.matmul(s, v)
+    ctx = ctx.permute(0,2,1,3).contiguous()
+
+    ctx.retain_grad()
+
+    return ctx
+
+class TestFMHA(unittest.TestCase):
+
+    def run_test(self, s, b):
+        print(f'Test s={s} b={b}')
+
+        torch.manual_seed(1234)
+        torch.cuda.manual_seed(1234)
+        
+        dtype = torch.float16
+        device = torch.device('cuda')
+
+        h = 16 
+        d = 64
+    
+        slens = [s] * b 
+        a = torch.tensor(np.array([0] + slens), dtype=torch.int32)
+        amask = torch.ones(b,h,s,s, dtype=dtype, device=device)
+        seqlens = torch.tensor(slens, dtype=torch.int32, device=device)
+        cu_seqlens = torch.cumsum(a, 0).to(dtype=torch.int32, device=device)
+        total = cu_seqlens[-1].item()
+    
+        qkv = torch.randn((b,s,h,3,d), device=device, dtype=dtype)
+    
+        qkv_vs = qkv.permute(0,1,3,2,4).contiguous().view(b*s, 3, h,d)
+    
+        qkv.requires_grad = True
+    
+        if b < 4:
+            ctx, S_ = mha.fwd_nl(qkv_vs, cu_seqlens, 0.0, s, True, None)
+        else:
+            ctx, S_ = mha.fwd(qkv_vs, cu_seqlens, 0.0, s, True, None)
+        ctx = ctx.view(b,s,h,d)
+    
+        ctx_ref = py_mha(qkv, amask, b,s,h,d)
+        self.assertTrue(torch.allclose(ctx_ref.float(), ctx.float(), atol=1e-3))
+    
+        labels = torch.randn_like(ctx_ref)
+        diff = ctx_ref - labels
+        l = (diff * diff).sum() / b
+        l.backward()
+    
+        dw = ctx_ref.grad.permute(0,2,1,3) 
+    
+        dw2 = dw.permute(0,2,1,3).clone().detach().contiguous()
+    
+        if b < 4:
+            dqkv2, _, _ = mha.bwd_nl(dw2, qkv_vs, S_, cu_seqlens, 0.0, s)
+        else:
+            dqkv2, _ = mha.bwd(dw2, qkv_vs, S_, cu_seqlens, 0.0, s)
+        
+        dqkv2 = dqkv2.permute(0,2,1,3).view(b,s, h,3,d)
+    
+        self.assertTrue(torch.allclose(qkv.grad.float(), dqkv2.float(), atol=1e-3))
+
+    def test_128(self):
+        self.run_test(128, 32)
+
+    def test_256(self):
+        self.run_test(256, 32)
+
+    def test_384(self):
+        self.run_test(384, 32)
+
+    def test_512(self):
+        self.run_test(512, 32)
+        self.run_test(512, 2)
+        self.run_test(512, 3)
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/fused_dense/test_fused_dense.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/fused_dense/test_fused_dense.py
new file mode 100644
index 0000000000..301ebf6b56
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/fused_dense/test_fused_dense.py
@@ -0,0 +1,44 @@
+import torch
+import unittest
+import torch.nn.functional as F
+from apex import fused_dense
+from torch import nn
+from apex import amp
+
+class FusedDenseTest(unittest.TestCase):
+    def setUp(self, seed=0):
+        torch.manual_seed(seed)
+        #torch.cuda.manual_seed_all(seed)
+
+        self.seq_length   = 512
+        self.sequences    = 3
+        self.hidden_dim   = 1024
+
+        self.ref_inputs = torch.randn(self.sequences*self.seq_length, self.hidden_dim,
+                                      dtype=torch.float16, device=torch.device("cuda")).int().half().requires_grad_(True)
+
+        self.tst_inputs = self.ref_inputs.clone().detach().requires_grad_(True)
+        self.dense = fused_dense.FusedDense(1024, 3072)
+        self.dense.half()
+        self.dense.cuda()
+
+
+    def test_fused_dense(self) :
+        y_tst = self.dense(self.tst_inputs)
+        y_ref = torch.matmul(self.ref_inputs,self.dense.weight.t())+self.dense.bias
+        dy = torch.randn_like(y_tst).half()
+        y_tst.backward(dy)
+        dw_ref = torch.matmul(dy.t(), self.ref_inputs)
+        dx_ref = torch.matmul(dy, self.dense.weight.clone())
+        db_ref = dy.sum(0, False)
+
+
+        self.assertTrue(torch.allclose(self.ref_inputs,  self.tst_inputs,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(y_ref,  y_tst,  atol=1e-3, rtol=1e-3, equal_nan=True))
+        self.assertTrue(torch.allclose(dw_ref, self.dense.weight.grad, atol=1e-3, rtol=1e-3, equal_nan=True))
+        self.assertTrue(torch.allclose(dx_ref, self.tst_inputs.grad, atol=1e-3, rtol=1e-3, equal_nan=True))
+        self.assertTrue(torch.allclose(db_ref, self.dense.bias.grad, atol=1e-3, rtol=1e-3, equal_nan=True))
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/layer_norm/test_fast_layer_norm.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/layer_norm/test_fast_layer_norm.py
new file mode 100644
index 0000000000..ca2f1e71ac
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/layer_norm/test_fast_layer_norm.py
@@ -0,0 +1,275 @@
+import unittest
+import sys
+import os
+
+import numpy as np
+import torch
+
+import fast_layer_norm as fln
+from apex.contrib.layer_norm.layer_norm import FastLayerNorm
+
+
+class GPUTimer:
+    def __init__(self, stream):
+        self.start_ = torch.cuda.Event(enable_timing=True)
+        self.stop_ = torch.cuda.Event(enable_timing=True)
+        self.stream_ = stream
+
+    def start(self):
+        self.stream_.record_event(self.start_)
+
+    def stop(self):
+        self.stream_.record_event(self.stop_)
+
+    def sync(self):
+        self.stream_.synchronize()
+
+    def millis(self):
+        return self.start_.elapsed_time(self.stop_)
+
+
+def size_in_bytes(t):
+    return torch.numel(t) * t.element_size()
+
+
+def metrics(y_ref, y, epsilon=1e-6):
+    y_ref = y_ref.float()
+    y = y.float()
+    relerr, mse = (
+        (y_ref - y).abs().sum() / (y_ref.abs().sum() + epsilon),
+        (y_ref - y).square().mean(),
+    )
+    return relerr.item(), mse.item()
+
+
+device = torch.device("cuda")
+fp32 = torch.float32
+fp16 = torch.float16
+bf16 = torch.bfloat16
+
+
+def backward_(dz, x, mu, rs, gamma):
+
+    wtype = gamma.dtype
+    itype = x.dtype
+    otype = dz.dtype
+    ctype = mu.dtype
+    mu = mu.unsqueeze(1)
+    rs = rs.unsqueeze(1)
+
+    hidden_size = gamma.numel()
+    y = rs * (x.to(ctype) - mu)
+    dbeta = dz.view(-1, hidden_size).sum(0, dtype=ctype)
+    dgamma = (dz * y).view(-1, hidden_size).sum(0, dtype=ctype)
+    dy = dz.view(-1, hidden_size).to(ctype) * gamma.unsqueeze(0).to(ctype)
+    mdy = dy.mean(1, keepdim=True, dtype=ctype)
+
+    mdyy = (dy * y).mean(1, keepdim=True, dtype=ctype)
+    dx = rs * (dy - mdyy * y - mdy)
+
+    return dx.to(itype), dgamma.to(wtype), dbeta.to(wtype)
+
+
+def benchmark_(S, B, hidden_size, itype, wtype, runs=100):
+    epsilon = 1e-5
+
+    x = torch.randn((S * B, hidden_size), dtype=itype, device=device)
+    beta = torch.randn(hidden_size, dtype=wtype, device=device)
+    gamma = torch.randn(hidden_size, dtype=wtype, device=device)
+    dz = torch.randn(x.shape, dtype=wtype, device=device)
+
+    stream = torch.cuda.Stream()
+    with torch.cuda.stream(stream):
+
+        timer = GPUTimer(stream)
+
+        # warmup
+        for r in range(runs):
+            z, mu, rsigma = fln.ln_fwd(x, gamma, beta, epsilon)
+
+        timer.start()
+        for r in range(runs):
+            z, mu, rsigma = fln.ln_fwd(x, gamma, beta, epsilon)
+        timer.stop()
+        timer.sync()
+
+        total_bytes_fwd = sum([size_in_bytes(t) for t in [x, z, gamma, beta, mu, rsigma]])
+
+        ms_fwd = timer.millis() / runs
+
+        print(
+            "[FWD] Time: {:.4f}ms Throughput: {:.4f} GB/sec".format(
+                ms_fwd, total_bytes_fwd * 1e-6 / ms_fwd
+            )
+        )
+
+        timer.start()
+        for r in range(runs):
+            dx, dgamma, dbeta, dbp, dgp = fln.ln_bwd(dz, x, mu, rsigma, gamma)
+        timer.stop()
+        timer.sync()
+
+        total_bytes_bwd = sum(
+            [
+                size_in_bytes(t)
+                for t in [dz, x, mu, rsigma, gamma, dx, dgamma, dbeta, dbp, dbp, dgp, dgp]
+            ]
+        )
+
+        ms_bwd = timer.millis() / runs
+
+        print(
+            "[BWD] Time: {:.4f}ms Throughput: {:.4f} GB/sec".format(
+                ms_bwd, total_bytes_bwd * 1e-6 / ms_bwd
+            )
+        )
+
+
+def test_(S, B, hidden_size, itype, wtype, ctype=fp32):
+
+    seed = 1243
+    torch.manual_seed(seed)
+    torch.cuda.manual_seed(seed)
+
+    otype = wtype
+    print("========================================================")
+    print(f"S={S} B={B} Hidden={hidden_size} {itype} {wtype}")
+    print("--------------------------------------------------------")
+
+    x = torch.randn(S * B, hidden_size, dtype=itype, device=device)
+    gamma = torch.randn(hidden_size, dtype=wtype, device=device) * 0.2
+    beta = torch.randn(hidden_size, dtype=wtype, device=device) * 0.2
+    epsilon = 1e-5
+
+    x.requires_grad = True
+    gamma.requires_grad = True
+    beta.requires_grad = True
+
+    mu_ref = x.mean(1, dtype=ctype, keepdim=True)
+    v = torch.square(x - mu_ref).mean(1, dtype=ctype, keepdim=True)
+    rs_ref = torch.rsqrt(v + epsilon)
+    y_ref = rs_ref * (x.to(ctype) - mu_ref)
+    z_ref = (gamma.unsqueeze(0) * (y_ref).to(otype) + beta.unsqueeze(0)).to(otype)
+
+    mu_ref = mu_ref.flatten()
+    rs_ref = rs_ref.flatten()
+
+    dz = torch.randn_like(z_ref)
+
+    # z_ref.backward(dz)
+    # dx_ref = x.grad
+    # dgamma_ref = gamma.grad
+    # dbeta_ref = beta.grad
+
+    dx_ref, dg_ref, db_ref = backward_(dz, x, mu_ref, rs_ref, gamma)
+
+    z, mu, rs = fln.ln_fwd(x, gamma, beta, epsilon)
+    dx, dg, db, dg_part, db_part = fln.ln_bwd(dz, x, mu, rs, gamma)
+
+    re_z, mse_z = metrics(z_ref, z)
+    re_mu, mse_mu = metrics(mu_ref, mu)
+    re_rs, mse_rs = metrics(rs_ref, rs)
+
+    re_dx, mse_dx = metrics(dx_ref, dx)
+    re_dg, mse_dg = metrics(dg_ref, dg)
+    re_db, mse_db = metrics(db_ref, db)
+
+    print(f" z: relerr={re_z :.4e} mse={mse_z :.4e}")
+    print(f"mu: relerr={re_mu:.4e} mse={mse_mu:.4e}")
+    print(f"rs: relerr={re_mu:.4e} mse={mse_mu:.4e}")
+
+    print(f"dx: relerr={re_dx:.4e} mse={mse_dx:.4e}")
+    print(f"dg: relerr={re_dg:.4e} mse={mse_dg:.4e}")
+    print(f"db: relerr={re_db:.4e} mse={mse_db:.4e}")
+
+    def check_err(x, relerr):
+        tol = 1e-3 if x.dtype == torch.float16 else 5e-6
+        return relerr < tol
+
+    return [
+        check_err(x, re)
+        for x, re in zip([z, mu, rs, dx, dg, db], [re_z, re_mu, re_rs, re_dx, re_dg, re_db])
+    ]
+
+
+class TestFastLayerNorm(unittest.TestCase):
+    def assertAll(self, l):
+        if not all(l):
+            print(l)
+        for x in l:
+            self.assertTrue(x)
+
+    def test_all_configs(self):
+
+        hidden_sizes = [
+            768,
+            1024,
+            1536,
+            2048,
+            2304,
+            3072,
+            3840,
+            4096,
+            5120,
+            6144,
+            8192,
+            10240,
+            12288,
+            12800,
+            15360,
+            16384,
+            18432,
+            20480,
+            24576,
+            25600,
+            30720,
+            32768,
+            40960,
+            49152,
+            65536,
+        ]
+
+        for h in hidden_sizes:
+            with self.subTest(f"hidden_size={h}"):
+                self.assertAll(test_(256, 2, h, fp32, fp32))
+                self.assertAll(test_(256, 2, h, fp16, fp16))
+                self.assertAll(test_(256, 2, h, fp32, fp16))
+                self.assertAll(test_(256, 2, h, bf16, bf16))
+                self.assertAll(test_(256, 2, h, fp32, bf16))
+
+    def test_run_benchmark(self):
+        for (S, B, hidden_size, runs) in (
+            (512, 32, 768, 1000),
+            (512, 32, 1024, 1000),
+            (512, 8, 4096, 1000),
+            (512, 8, 5120, 1000),
+            (512, 8, 6144, 1000),
+            (256, 2, 20480, 500),
+            (256, 2, 25600, 500),
+            (256, 2, 40960, 250),
+            (256, 2, 65536, 250),
+        ):
+            with self.subTest(f"(S, B, hidden_size)=({S}, {B}, {hidden_size})"):
+                benchmark_(S, B, hidden_size, fp16, fp16, runs)
+
+    def test_compat_with_autocast(self):
+        autocast_dtypes = (
+            (torch.half, torch.bfloat16) if torch.cuda.is_bf16_supported() else (torch.half,)
+        )
+        input_shape = (512, 32, 768)
+        layer_norm = FastLayerNorm(input_shape[-1]).cuda()
+        input = torch.randn(input_shape).cuda()
+
+        for dtype in autocast_dtypes:
+            layer_norm.zero_grad(set_to_none=True)
+            with self.subTest(f"autocast_dtype={dtype}"):
+                with torch.cuda.amp.autocast(enabled=True, dtype=dtype):
+                    out = layer_norm(input)
+                    self.assertEqual(dtype, out.dtype)
+                grad = torch.randn_like(out)
+                out.backward(grad)
+                self.assertEqual(torch.float32, layer_norm.weight.grad.dtype)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_encdec_multihead_attn.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_encdec_multihead_attn.py
new file mode 100644
index 0000000000..26c723824d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_encdec_multihead_attn.py
@@ -0,0 +1,136 @@
+import torch
+
+import unittest
+
+from apex.contrib.multihead_attn import EncdecMultiheadAttn
+
+class EncdecMultiheadAttnTest(unittest.TestCase):
+    def setUp(self, seed=1234):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+        self.seq_length   = 80
+        self.sequences    = 10
+        self.hidden_dim   = 1024
+        self.heads        = 16
+        self.dropout_prob = 0.0
+
+        self.ref_layer = EncdecMultiheadAttn(self.hidden_dim, 
+                                             self.heads, 
+                                             dropout=self.dropout_prob, 
+                                             bias=False, 
+                                             include_norm_add=False, 
+                                             impl='default')
+        self.ref_layer.cuda().half()
+        self.ref_layer.reset_parameters()
+        self.ref_inputs_q = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                        dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+        self.ref_inputs_k = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                        dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+
+        # Reset seed so parameters are identical
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        
+        self.tst_layer = EncdecMultiheadAttn(self.hidden_dim, 
+                                             self.heads, 
+                                             dropout=self.dropout_prob, 
+                                             bias=False, 
+                                             include_norm_add=False, 
+                                             impl='fast')
+        self.tst_layer.cuda().half()
+        self.tst_layer.reset_parameters()
+        
+        self.tst_inputs_q = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                        dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+        self.tst_inputs_k = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                        dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+
+    def test_encdec_multihead_attn(self) :
+        grads         = torch.randn_like(self.tst_inputs_q)
+
+        ref_outputs,_ = self.ref_layer.forward(self.ref_inputs_q, 
+                                               self.ref_inputs_k, 
+                                               self.ref_inputs_k,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+
+        tst_outputs,_ = self.tst_layer.forward(self.tst_inputs_q, 
+                                               self.tst_inputs_k, 
+                                               self.tst_inputs_k,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+        
+        self.ref_inputs_q.backward(grads)
+        self.tst_inputs_q.backward(grads)
+
+        self.assertTrue(torch.allclose(self.ref_inputs_q,  self.tst_inputs_q,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(self.ref_inputs_k,  self.tst_inputs_k,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3))
+    
+    def test_encdec_multihead_attn_time_mask(self) :
+        grads          = torch.randn_like(self.tst_inputs_q)
+        time_mask_byte = torch.triu(torch.ones(self.tst_inputs_q.size(0), self.tst_inputs_k.size(0), device=torch.device("cuda"), dtype=torch.uint8), 1)
+        time_mask_bool = time_mask_byte.to(torch.bool)
+        
+        ref_outputs,_ = self.ref_layer.forward(self.ref_inputs_q, 
+                                               self.ref_inputs_k, 
+                                               self.ref_inputs_k,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=time_mask_bool,
+                                               is_training=True)
+
+        tst_outputs,_ = self.tst_layer.forward(self.tst_inputs_q, 
+                                               self.tst_inputs_k, 
+                                               self.tst_inputs_k,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=time_mask_byte,
+                                               is_training=True)
+        
+        self.ref_inputs_q.backward(grads)
+        self.tst_inputs_q.backward(grads)
+
+        self.assertTrue(torch.allclose(self.ref_inputs_q,  self.tst_inputs_q,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(self.ref_inputs_k,  self.tst_inputs_k,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3))
+    
+    def test_encdec_multihead_attn_pad_mask(self) :
+        grads         = torch.randn_like(self.tst_inputs_q)
+        pad_mask_byte = torch.tril(torch.ones(self.tst_inputs_k.size(1), self.tst_inputs_k.size(0), device=torch.device("cuda"), dtype=torch.uint8), 1)
+        pad_mask_bool = pad_mask_byte.to(torch.bool)
+        
+        ref_outputs,_ = self.ref_layer.forward(self.ref_inputs_q, 
+                                               self.ref_inputs_k, 
+                                               self.ref_inputs_k,
+                                               key_padding_mask=pad_mask_bool, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+
+        tst_outputs,_ = self.tst_layer.forward(self.tst_inputs_q, 
+                                               self.tst_inputs_k, 
+                                               self.tst_inputs_k,
+                                               key_padding_mask=pad_mask_byte, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+        
+        self.ref_inputs_q.backward(grads)
+        self.tst_inputs_q.backward(grads)
+
+        self.assertTrue(torch.allclose(self.ref_inputs_q,  self.tst_inputs_q,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(self.ref_inputs_k,  self.tst_inputs_k,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3))
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_encdec_multihead_attn_norm_add.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_encdec_multihead_attn_norm_add.py
new file mode 100644
index 0000000000..1313339e94
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_encdec_multihead_attn_norm_add.py
@@ -0,0 +1,77 @@
+import torch
+
+import unittest
+
+from apex.contrib.multihead_attn import EncdecMultiheadAttn
+
+class EncdecMultiheadAttnNormAddTest(unittest.TestCase):
+    def setUp(self, seed=1234):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+        self.seq_length   = 80
+        self.sequences    = 10
+        self.hidden_dim   = 1024
+        self.heads        = 16
+        self.dropout_prob = 0.0
+
+        self.ref_layer = EncdecMultiheadAttn(self.hidden_dim, 
+                                             self.heads, 
+                                             dropout=self.dropout_prob, 
+                                             bias=False, 
+                                             include_norm_add=True, 
+                                             impl='default')
+        self.ref_layer.cuda().half()
+        self.ref_layer.reset_parameters()
+        self.ref_inputs_q = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                        dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+        self.ref_inputs_k = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                        dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+
+        # Reset seed so parameters are identical
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        
+        self.tst_layer = EncdecMultiheadAttn(self.hidden_dim, 
+                                             self.heads, 
+                                             dropout=self.dropout_prob, 
+                                             bias=False, 
+                                             include_norm_add=True, 
+                                             impl='fast')
+        self.tst_layer.cuda().half()
+        self.tst_layer.reset_parameters()
+        
+        self.tst_inputs_q = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                        dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+        self.tst_inputs_k = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                        dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+
+    def test_encdec_multihead_attn_norm_add(self) :
+        grads         = torch.randn_like(self.tst_inputs_q)
+
+        ref_outputs,_ = self.ref_layer.forward(self.ref_inputs_q, 
+                                               self.ref_inputs_k, 
+                                               self.ref_inputs_k,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+
+        tst_outputs,_ = self.tst_layer.forward(self.tst_inputs_q, 
+                                               self.tst_inputs_k, 
+                                               self.tst_inputs_k,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+        
+        self.ref_inputs_q.backward(grads)
+        self.tst_inputs_q.backward(grads)
+
+        self.assertTrue(torch.allclose(self.ref_inputs_q,  self.tst_inputs_q,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(self.ref_inputs_k,  self.tst_inputs_k,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(self.ref_inputs_q.grad, self.tst_inputs_q.grad, atol=1e-3, rtol=1e-3))
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_fast_self_multihead_attn_bias.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_fast_self_multihead_attn_bias.py
new file mode 100644
index 0000000000..b4bbf342b6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_fast_self_multihead_attn_bias.py
@@ -0,0 +1,77 @@
+import torch
+
+import unittest
+
+from apex.contrib.multihead_attn import SelfMultiheadAttn
+
+class SelfMultiheadAttnTest(unittest.TestCase):
+    def setUp(self, seed=1234):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+        self.seq_length   = 80
+        self.sequences    = 10
+        self.hidden_dim   = 1024
+        self.heads        = 16
+        self.dropout_prob = 0.0
+
+        self.ref_layer = SelfMultiheadAttn(self.hidden_dim, 
+                                           self.heads, 
+                                           dropout=self.dropout_prob, 
+                                           bias=True, 
+                                           include_norm_add=False, 
+                                           separate_qkv_params=True, 
+                                           mask_additive=True, 
+                                           impl='default')
+        self.ref_layer.cuda().half()
+        self.ref_layer.reset_parameters()
+        self.ref_inputs = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                      dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+        # Reset seed so parameters are identical
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        
+        self.tst_layer = SelfMultiheadAttn(self.hidden_dim, 
+                                           self.heads, 
+                                           dropout=self.dropout_prob, 
+                                           bias=True, 
+                                           include_norm_add=False, 
+                                           separate_qkv_params=True, 
+                                           mask_additive=True, 
+                                           impl='fast')
+        self.tst_layer.cuda().half()
+        self.tst_layer.reset_parameters()
+        
+        self.tst_inputs = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                      dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+    
+    def test_self_multihead_attn_additive_mask(self) :
+        grads         = torch.randn_like(self.tst_inputs)
+        mask = ((torch.randn(self.sequences, self.seq_length) > 0) * -10000.0).half().cuda()
+
+        ref_outputs,_ = self.ref_layer.forward(self.ref_inputs, 
+                                               self.ref_inputs, 
+                                               self.ref_inputs,
+                                               key_padding_mask=mask, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+
+        tst_outputs,_ = self.tst_layer.forward(self.tst_inputs, 
+                                               self.tst_inputs, 
+                                               self.tst_inputs,
+                                               key_padding_mask=mask, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+
+        
+        self.ref_inputs.backward(grads)
+        self.tst_inputs.backward(grads)
+
+        self.assertTrue(torch.allclose(self.ref_inputs,  self.tst_inputs,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3))
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_mha_fused_softmax.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_mha_fused_softmax.py
new file mode 100644
index 0000000000..60d9541847
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_mha_fused_softmax.py
@@ -0,0 +1,42 @@
+import torch
+import unittest
+import torch.nn.functional as F
+from apex.contrib.multihead_attn import fast_mask_softmax_dropout_func
+
+class FusedSoftmaxTest(unittest.TestCase):
+    def setUp(self, seed=1234):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+        self.seq_length   = 80
+        self.sequences    = 10
+        self.hidden_dim   = 1024
+        self.heads        = 16
+        self.dropout_prob = 0.0
+
+        self.mask = (torch.randn(self.sequences,self.seq_length)>0).cuda()
+        self.mask = self.mask.half()*-10000
+        self.ref_inputs = torch.randn(self.heads * self.sequences, self.seq_length, self.seq_length, 
+                                      dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+        
+        self.tst_inputs = self.ref_inputs.clone().detach().requires_grad_(True)
+
+    def test_fused_softmax(self) :
+        grads = torch.randn_like(self.tst_inputs)
+        y_ref = self.ref_inputs.view(self.sequences, self.heads, self.seq_length, self.seq_length)
+        y_ref = y_ref + self.mask.unsqueeze(1).unsqueeze(2)
+        y_ref = y_ref.view(self.sequences*self.heads, self.seq_length, self.seq_length) 
+        y_ref = F.softmax(y_ref, dim=-1)
+        y_ref = torch._fused_dropout(y_ref, 1.0)    
+   
+        y_tst = fast_mask_softmax_dropout_func(True, self.heads, self.tst_inputs, self.mask, True, 0.0)        
+        y_ref[0].backward(grads)
+        y_tst.backward(grads)
+
+        self.assertTrue(torch.allclose(self.ref_inputs,  self.tst_inputs,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(y_ref[0], y_tst, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3))
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_self_multihead_attn.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_self_multihead_attn.py
new file mode 100644
index 0000000000..7438e177ed
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_self_multihead_attn.py
@@ -0,0 +1,130 @@
+import torch
+
+import unittest
+
+from apex.contrib.multihead_attn import SelfMultiheadAttn
+
+class SelfMultiheadAttnTest(unittest.TestCase):
+    def setUp(self, seed=1234):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+        self.seq_length   = 80
+        self.sequences    = 10
+        self.hidden_dim   = 1024
+        self.heads        = 16
+        self.dropout_prob = 0.0
+
+        self.ref_layer = SelfMultiheadAttn(self.hidden_dim, 
+                                           self.heads, 
+                                           dropout=self.dropout_prob, 
+                                           bias=False, 
+                                           include_norm_add=False, 
+                                           impl='default')
+        self.ref_layer.cuda().half()
+        self.ref_layer.reset_parameters()
+        self.ref_inputs = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                      dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+
+        # Reset seed so parameters are identical
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        
+        self.tst_layer = SelfMultiheadAttn(self.hidden_dim, 
+                                           self.heads, 
+                                           dropout=self.dropout_prob, 
+                                           bias=False, 
+                                           include_norm_add=False, 
+                                           impl='fast')
+        self.tst_layer.cuda().half()
+        self.tst_layer.reset_parameters()
+        
+        self.tst_inputs = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                      dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+
+    def test_self_multihead_attn(self) :
+        grads         = torch.randn_like(self.tst_inputs)
+
+        ref_outputs,_ = self.ref_layer.forward(self.ref_inputs, 
+                                               self.ref_inputs, 
+                                               self.ref_inputs,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+
+        tst_outputs,_ = self.tst_layer.forward(self.tst_inputs, 
+                                               self.tst_inputs, 
+                                               self.tst_inputs,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+        
+        self.ref_inputs.backward(grads)
+        self.tst_inputs.backward(grads)
+
+        self.assertTrue(torch.allclose(self.ref_inputs,  self.tst_inputs,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3))
+
+    def test_self_multihead_attn_time_mask(self) :
+        grads         = torch.randn_like(self.tst_inputs)
+        time_mask_byte= torch.triu(torch.ones(self.tst_inputs.size(0), self.tst_inputs.size(0), device=torch.device("cuda"), dtype=torch.uint8), 1)
+        time_mask_bool= time_mask_byte.to(torch.bool)
+
+        ref_outputs,_ = self.ref_layer.forward(self.ref_inputs, 
+                                               self.ref_inputs, 
+                                               self.ref_inputs,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=time_mask_bool,
+                                               is_training=True)
+
+        tst_outputs,_ = self.tst_layer.forward(self.tst_inputs, 
+                                               self.tst_inputs, 
+                                               self.tst_inputs,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=time_mask_byte,
+                                               is_training=True)
+
+        
+        self.ref_inputs.backward(grads)
+        self.tst_inputs.backward(grads)
+
+        self.assertTrue(torch.allclose(self.ref_inputs,  self.tst_inputs,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3))
+    
+    def test_self_multihead_attn_pad_mask(self) :
+        grads         = torch.randn_like(self.tst_inputs)
+        pad_mask_byte = torch.tril(torch.ones(self.tst_inputs.size(1), self.tst_inputs.size(0), device=torch.device("cuda"), dtype=torch.uint8), 1)
+        pad_mask_bool = pad_mask_byte.to(torch.bool)
+
+        ref_outputs,_ = self.ref_layer.forward(self.ref_inputs, 
+                                               self.ref_inputs, 
+                                               self.ref_inputs,
+                                               key_padding_mask=pad_mask_bool, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+
+        tst_outputs,_ = self.tst_layer.forward(self.tst_inputs, 
+                                               self.tst_inputs, 
+                                               self.tst_inputs,
+                                               key_padding_mask=pad_mask_byte, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+
+        
+        self.ref_inputs.backward(grads)
+        self.tst_inputs.backward(grads)
+
+        self.assertTrue(torch.allclose(self.ref_inputs,  self.tst_inputs,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3))
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_self_multihead_attn_norm_add.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_self_multihead_attn_norm_add.py
new file mode 100644
index 0000000000..efbb7f200d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/multihead_attn/test_self_multihead_attn_norm_add.py
@@ -0,0 +1,72 @@
+import torch
+
+import unittest
+
+from apex.contrib.multihead_attn import SelfMultiheadAttn
+
+class SelfMultiheadAttnNormAddTest(unittest.TestCase):
+    def setUp(self, seed=1234):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+        self.seq_length   = 80
+        self.sequences    = 10
+        self.hidden_dim   = 1024
+        self.heads        = 16
+        self.dropout_prob = 0.0
+
+        self.ref_layer = SelfMultiheadAttn(self.hidden_dim, 
+                                           self.heads, 
+                                           dropout=self.dropout_prob, 
+                                           bias=False, 
+                                           include_norm_add=True, 
+                                           impl='default')
+        self.ref_layer.cuda().half()
+        self.ref_layer.reset_parameters()
+        self.ref_inputs = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                      dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+
+        # Reset seed so parameters are identical
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+        
+        self.tst_layer = SelfMultiheadAttn(self.hidden_dim, 
+                                           self.heads, 
+                                           dropout=self.dropout_prob, 
+                                           bias=False, 
+                                           include_norm_add=True, 
+                                           impl='fast')
+        self.tst_layer.cuda().half()
+        self.tst_layer.reset_parameters()
+        
+        self.tst_inputs = torch.randn(self.seq_length, self.sequences, self.hidden_dim, 
+                                      dtype=torch.float16, device=torch.device("cuda")).requires_grad_(True)
+
+    def test_self_multihead_attn_norm_add(self) :
+        grads         = torch.randn_like(self.tst_inputs)
+
+        ref_outputs,_ = self.ref_layer.forward(self.ref_inputs, 
+                                               self.ref_inputs, 
+                                               self.ref_inputs,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+
+        tst_outputs,_ = self.tst_layer.forward(self.tst_inputs, 
+                                               self.tst_inputs, 
+                                               self.tst_inputs,
+                                               key_padding_mask=None, 
+                                               need_weights=False, 
+                                               attn_mask=None,
+                                               is_training=True)
+        
+        self.ref_inputs.backward(grads)
+        self.tst_inputs.backward(grads)
+
+        self.assertTrue(torch.allclose(self.ref_inputs,  self.tst_inputs,  atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(ref_outputs, tst_outputs, atol=1e-3, rtol=1e-3))
+        self.assertTrue(torch.allclose(self.ref_inputs.grad, self.tst_inputs.grad, atol=1e-3, rtol=1e-3))
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/test_label_smoothing.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/test_label_smoothing.py
new file mode 100644
index 0000000000..70e9f3d099
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/test_label_smoothing.py
@@ -0,0 +1,128 @@
+import torch
+from apex.contrib import xentropy as label_smoothing
+import unittest
+
+import warnings
+import random
+import numpy as np
+import time
+
+def label_smoothing_raw(x, target, padding_idx, smoothing):
+    logprobs = torch.nn.functional.log_softmax(x, dim=-1, dtype=torch.float32)
+
+    non_pad_mask = (target != padding_idx)
+    nll_loss = -logprobs.gather(dim=-1, index=target.unsqueeze(1))
+    nll_loss = nll_loss.squeeze(1)[non_pad_mask]
+    smooth_loss = -logprobs.mean(dim=-1)[non_pad_mask]
+    loss = (1.0 - smoothing) * nll_loss + smoothing * smooth_loss
+    return loss
+
+def label_smoothing_opt_1(x, target, padding_idx, smoothing):
+    logprobs = torch.nn.functional.log_softmax(x, dim=-1, dtype=torch.float32)
+
+    pad_mask = (target == padding_idx)
+    ll_loss = logprobs.gather(dim=-1, index=target.unsqueeze(1)).squeeze(1)
+    smooth_loss = logprobs.mean(dim=-1)
+    loss = (smoothing - 1.0) * ll_loss - smoothing * smooth_loss
+    loss.masked_fill_(pad_mask, 0)
+    return loss
+
+class LabelSmoothingTest(unittest.TestCase):
+    def setUp(self, seed=1234):
+        random.seed(seed)
+        np.random.seed(seed)
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+        # Set pytorch print precision
+        torch.set_printoptions(precision=10)
+
+    def gen_test_inputs(self, N, T, H, smoothing, padding_idx):
+        logits = torch.randn((N*T, H), dtype=torch.half, device='cuda',
+            requires_grad=True)
+        labels = torch.randint(0, H, [N*T], device='cuda')
+        for i in random.sample(range(N*T), N*T//6):
+            labels[i] = padding_idx
+        half_to_float = (logits.dtype == torch.half)
+
+        return logits, labels, half_to_float
+
+    def print_max_diff_elem(self, ref, tst):
+        ref, tst = ref.flatten(), tst.flatten()
+        diff = (ref - tst).abs().max()
+        idx = (ref - tst).abs().argmax()
+        print("Max atol idx: {}, diff: {:.6f}, ref: {:.6f}, tst: {:.6f}".format(
+            idx, diff, ref[idx], tst[idx]))
+
+    def test_label_smoothing_function(self):
+        # Set label smoothing configuration
+        smoothing, padding_idx = 0.1, 0
+        N, T, H = 128, 74, 32320
+        iters = 10
+        loss_func = label_smoothing.SoftmaxCrossEntropyLoss.apply
+
+        for i in range(iters):
+            logits, labels, half_to_float = self.gen_test_inputs(
+                N, T, H, smoothing, padding_idx)
+    
+            # Run original softmax cross entropy with label smoothing
+            logits.grad = None
+            losses = label_smoothing_raw(logits, labels, padding_idx, smoothing)
+            loss = losses.sum()
+            loss.backward()
+            
+            ref_loss = loss.clone().detach()
+            ref_grad = logits.grad.clone().detach()
+
+            # Run optimized softmax cross entropy with label smoothing
+            logits.grad = None
+            losses = loss_func(logits, labels, smoothing, padding_idx, half_to_float)
+            loss = losses.sum()
+            loss.backward()
+
+            val_loss = loss.clone().detach()
+            val_grad = logits.grad.clone().detach()
+
+            # Validate
+            self.print_max_diff_elem(ref_grad, val_grad)
+            self.assertTrue(torch.allclose(ref_loss, val_loss, atol=1e-5, rtol=1e-5))
+            self.assertTrue(torch.allclose(ref_grad, val_grad, atol=1e-5, rtol=1e-5))
+
+    def test_label_smoothing_perf(self):
+        # Set label smoothing configuration
+        smoothing, padding_idx = 0.1, 0
+        N, T, H = 128, 74, 32320
+        iters = 1000
+        loss_func = label_smoothing.SoftmaxCrossEntropyLoss.apply
+        print()
+
+        logits, labels, half_to_float = self.gen_test_inputs(
+            N, T, H, smoothing, padding_idx)
+    
+        # Run original softmax cross entropy with label smoothing
+        torch.cuda.synchronize()
+        ts = time.time()
+        for i in range(iters):
+            logits.grad = None
+            losses = label_smoothing_raw(logits, labels, padding_idx, smoothing)
+            loss = losses.sum() / N
+            loss.backward()
+        torch.cuda.synchronize()
+        print("Raw time {:.2f} s elapsed for {} iterations, norm {:.4f}".format(
+            time.time() - ts, iters, logits.grad.norm()))
+            
+        # Run optimized softmax cross entropy with label smoothing
+        torch.cuda.synchronize()
+        ts = time.time()
+        for i in range(iters):
+            logits.grad = None
+            losses = loss_func(logits, labels, smoothing, padding_idx, half_to_float)
+            loss = losses.sum() / N
+            loss.backward()
+        torch.cuda.synchronize()
+        print("Opt time {:.2f} s elapsed for {} iterations, norm {:.4f}".format(
+            time.time() - ts, iters, logits.grad.norm()))
+
+if __name__ == '__main__':
+    unittest.main()
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/test_transducer_joint.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/test_transducer_joint.py
new file mode 100644
index 0000000000..aba8624d95
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/test_transducer_joint.py
@@ -0,0 +1,157 @@
+import torch
+import unittest
+from apex.contrib.transducer import TransducerJoint
+import transducer_ref
+
+class TransducerJointTest(unittest.TestCase):
+    def setUp(self, seed=1234):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+    def gen_input(self, for_vector_kernel):
+        self.B = 4
+        T_min = 51
+        T_max = 101
+        U_min = 12
+        U_max = 25
+        if for_vector_kernel:
+            H = 512
+        else:
+            H = 509
+        dtype = torch.float16
+        device = "cuda"
+
+        self.f_tst = torch.randn((self.B, T_max, H), dtype=dtype, requires_grad=True, device=device)
+        self.g_tst = torch.randn((self.B, U_max, H), dtype=dtype, requires_grad=True, device=device)
+        self.h_grad = torch.randn(self.B, T_max, U_max, H, dtype=dtype, device=device)
+        self.f_len = torch.randint(T_min, T_max+1, (self.B,), dtype=torch.int, device=device) 
+        self.g_len = torch.randint(U_min, U_max+1, (self.B,), dtype=torch.int, device=device)
+        self.f_len[torch.randint(0, self.B, (1,)).item()] = T_max
+        self.g_len[torch.randint(0, self.B, (1,)).item()] = U_max
+        self.dropout_prob = 0.5
+
+        # Make sure gradients from out-of-bound locations are zero. This should be guaranteed by 
+        # the loss function
+        for b in range(self.B):
+            self.h_grad[b, self.f_len[b]:, :, :] = 0
+            self.h_grad[b, :, self.g_len[b]:, :] = 0
+        self.h_grad_packed = self._pack(self.h_grad, self.f_len, self.g_len)
+        
+
+    def _pack(self, x, f_len, g_len):
+        B = x.size(0)
+        list_x = []
+        for b in range(B):
+            list_x_row = [x[b, t, :g_len[b]] for t in range(f_len[b])]
+            x_row = torch.cat(list_x_row)
+            list_x.append(x_row)
+        x_packed = torch.cat(list_x).data.clone()
+        x_packed.requires_grad = True
+        batch_offset = torch.cumsum(f_len * g_len, dim=0)
+        return x_packed
+
+    def _unpack(self, x, f_len, g_len):
+        batch_offset = torch.cumsum(f_len * g_len, dim=0)
+        x_unpacked = torch.zeros_like(self.h_grad, dtype=torch.uint8)
+        B = self.h_grad.size(0)
+        H = self.h_grad.size(-1)
+        for b in range(B):
+            my_batch_offset = 0 if b == 0 else batch_offset[b-1]
+            my_f_len = f_len[b]
+            my_g_len = g_len[b]
+            for t in range(my_f_len):
+                x_unpacked[b, t, :my_g_len] = x[my_batch_offset + t*my_g_len : 
+                                                my_batch_offset + t*my_g_len + my_g_len]
+        return x_unpacked
+        
+    def run_transducer_joint(self, for_vector_kernel, pack_output, relu, dropout):
+        self.gen_input(for_vector_kernel=for_vector_kernel)
+        # Generate reference
+        f_ref = self.f_tst.data.clone()
+        g_ref = self.g_tst.data.clone()
+        f_ref.requires_grad = True
+        g_ref.requires_grad = True
+        
+        my_joint = TransducerJoint(pack_output=pack_output, relu=relu, dropout=dropout, 
+                                    dropout_prob=self.dropout_prob, probe_mask=True)
+        if not pack_output:
+            h_tst = my_joint(   f=self.f_tst, 
+                                g=self.g_tst, 
+                                f_len=self.f_len, 
+                                g_len=self.g_len)
+            h_tst.backward(self.h_grad)
+            if dropout:
+                mask = my_joint.mask_probe[0]
+        else:
+            batch_offset = torch.cumsum(self.f_len * self.g_len, dim=0)
+            h_tst = my_joint(   f=self.f_tst, 
+                                g=self.g_tst, 
+                                f_len=self.f_len, 
+                                g_len=self.g_len, 
+                                batch_offset=batch_offset, 
+                                packed_batch=batch_offset[-1])
+            h_tst.backward(self.h_grad_packed)
+            if dropout:
+                mask_packed = my_joint.mask_probe[0]
+                mask = self._unpack(mask_packed, self.f_len, self.g_len)
+
+        # reference
+        h_ref, f_grad_ref, g_grad_ref \
+            = transducer_ref.transducer_joint_reference(f=f_ref, 
+                                                        g=g_ref, 
+                                                        h_grad=self.h_grad, 
+                                                        f_len=self.f_len, 
+                                                        g_len=self.g_len, 
+                                                        pack_output=pack_output,
+                                                        relu=relu,
+                                                        dropout=dropout,
+                                                        dropout_prob=self.dropout_prob,
+                                                        mask=mask if dropout else None)
+        
+        f_grad_tst = self.f_tst.grad
+        g_grad_tst = self.g_tst.grad
+        
+        self.assertTrue(torch.allclose(h_ref, h_tst, atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(f_grad_ref, f_grad_tst, atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(g_grad_ref, g_grad_tst, atol=1e-4, rtol=1e-4))
+
+    def test_transducer_joint(self):
+        self.run_transducer_joint(for_vector_kernel=True, pack_output=True, relu=False, dropout=False)
+
+    def test_transducer_joint_vec(self):
+        self.run_transducer_joint(for_vector_kernel=True, pack_output=False, relu=False, dropout=False)
+
+    def test_transducer_joint_pack(self):
+        self.run_transducer_joint(for_vector_kernel=False, pack_output=True, relu=False, dropout=False)
+
+    def test_transducer_joint_vec_pack(self):
+        self.run_transducer_joint(for_vector_kernel=True, pack_output=True, relu=False, dropout=False)
+
+    def test_transducer_joint_relu(self):
+        self.run_transducer_joint(for_vector_kernel=True, pack_output=True, relu=True, dropout=False)
+
+    def test_transducer_joint_vec_relu(self):
+        self.run_transducer_joint(for_vector_kernel=True, pack_output=False, relu=True, dropout=False)
+
+    def test_transducer_joint_pack_relu(self):
+        self.run_transducer_joint(for_vector_kernel=False, pack_output=True, relu=True, dropout=False)
+
+    def test_transducer_joint_vec_pack_relu(self):
+        self.run_transducer_joint(for_vector_kernel=True, pack_output=True, relu=True, dropout=False)
+
+    def test_transducer_joint_relu_dropout(self):
+        self.run_transducer_joint(for_vector_kernel=True, pack_output=True, relu=True, dropout=True)
+
+    def test_transducer_joint_vec_relu_dropout(self):
+        self.run_transducer_joint(for_vector_kernel=True, pack_output=False, relu=True, dropout=True)
+
+    def test_transducer_joint_pack_relu_dropout(self):
+        self.run_transducer_joint(for_vector_kernel=False, pack_output=True, relu=True, dropout=True)
+
+    def test_transducer_joint_vec_pack_relu_dropout(self):
+        self.run_transducer_joint(for_vector_kernel=True, pack_output=True, relu=True, dropout=True)
+
+
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/test_transducer_loss.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/test_transducer_loss.py
new file mode 100644
index 0000000000..0f64ab5f9c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/test_transducer_loss.py
@@ -0,0 +1,133 @@
+import torch
+import unittest
+from apex.contrib.transducer import TransducerLoss
+import transducer_ref
+
+class TransducerLossTest(unittest.TestCase):
+    def setUp(self, seed=1234):
+        torch.manual_seed(seed)
+        torch.cuda.manual_seed_all(seed)
+
+    def gen_input(self, scalar_t, for_vector_kernel):
+        self.B = 5
+        T_min = 23
+        T_max = 51
+        U_min = 12
+        U_max = 25
+        V = 16 if for_vector_kernel else 14
+        self.blank_idx = V - 1
+        device = "cuda"
+
+        self.x_tst = torch.randn((self.B, T_max, U_max, V), dtype=scalar_t, requires_grad=True, 
+                                    device=device)
+        self.y = torch.randint(0, self.blank_idx, (self.B, U_max-1), dtype=torch.int, device=device)
+        self.f_len = torch.randint(T_min, T_max+1, (self.B,), dtype=torch.int, device=device) 
+        self.y_len = torch.randint(U_min-1, U_max, (self.B,), dtype=torch.int, device=device)
+        self.f_len[torch.randint(0, self.B, (1,)).item()] = T_max
+        self.y_len[torch.randint(0, self.B, (1,)).item()] = U_max-1
+        self.x_tst_packed, self.batch_offset = self._pack(self.x_tst)
+        # Generate reference
+        x_ref = self.x_tst.data.clone()
+        x_ref.requires_grad = True
+        loss_grad = torch.ones(x_ref.size(0), dtype=x_ref.dtype, device=x_ref.device)/x_ref.size(0)
+        _, _, self.grad_ref, self.loss_ref \
+            = transducer_ref.transducer_loss_reference( x=x_ref, 
+                                                        label=self.y, 
+                                                        f_len=self.f_len, 
+                                                        y_len=self.y_len, 
+                                                        blank_idx=self.blank_idx, 
+                                                        loss_grad=loss_grad)
+
+    def _pack(self, x):
+        list_x = []
+        for b in range(self.B):
+            list_x_row = [x[b, t, : self.y_len[b]+1] for t in range(self.f_len[b])]
+            x_row = torch.cat(list_x_row)
+            list_x.append(x_row)
+        x_packed = torch.cat(list_x).data.clone()
+        x_packed.requires_grad = True
+        batch_offset = torch.cumsum(self.f_len * (self.y_len+1), dim=0)
+        return x_packed, batch_offset
+
+    def _unpack(self, x):
+        x_unpacked = torch.zeros(self.B, self.f_len.max(), self.y_len.max()+1, x.size(-1), 
+                                    dtype=x.dtype, device=x.device)
+        for b in range(self.B):
+            my_batch_offset = 0 if b == 0 else self.batch_offset[b-1]
+            my_f_len = self.f_len[b]
+            my_g_len = self.y_len[b] + 1
+            for t in range(my_f_len):
+                for u in range(my_g_len):
+                    x_unpacked[b, t, u] = x[my_batch_offset + t*my_g_len + u]
+        return x_unpacked
+
+    def run_transducer_loss(self, scalar_t, fuse_softmax_backward, packed_input, for_vector_kernel):
+        self.gen_input(scalar_t, for_vector_kernel)
+        my_loss = TransducerLoss(  fuse_softmax_backward=fuse_softmax_backward, 
+                                    packed_input=packed_input) 
+        if not packed_input:
+            loss_tst = my_loss( x=self.x_tst,
+                                label=self.y, 
+                                f_len=self.f_len, 
+                                y_len=self.y_len, 
+                                blank_idx=self.blank_idx)
+            loss_tst.mean().backward() 
+            grad_tst = self.x_tst.grad
+        else:
+            loss_tst = my_loss( x=self.x_tst_packed,
+                                label=self.y, 
+                                f_len=self.f_len, 
+                                y_len=self.y_len, 
+                                blank_idx=self.blank_idx,
+                                batch_offset=self.batch_offset, 
+                                max_f_len=max(self.f_len))
+            loss_tst.mean().backward()
+            grad_tst_packed = self.x_tst_packed.grad
+            grad_tst = self._unpack(grad_tst_packed)
+        
+        return loss_tst, grad_tst
+
+    def test_transducer_loss_fp32(self):
+        loss_tst, grad_tst = self.run_transducer_loss(  scalar_t=torch.float32,
+                                                        fuse_softmax_backward=False,
+                                                        packed_input=False,
+                                                        for_vector_kernel=False)
+        self.assertTrue(torch.allclose(self.loss_ref, loss_tst, atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(self.grad_ref, grad_tst, atol=1e-5, rtol=1e-5))
+
+    def test_transducer_loss_fp16(self):
+        loss_tst, grad_tst = self.run_transducer_loss(  scalar_t=torch.float16,
+                                                        fuse_softmax_backward=False,
+                                                        packed_input=False,
+                                                        for_vector_kernel=False)
+        self.assertTrue(torch.allclose(self.loss_ref, loss_tst, atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(self.grad_ref, grad_tst, atol=1e-4, rtol=1e-3))
+
+    def test_transducer_loss_fp16_backward_fusion(self):
+        loss_tst, grad_tst = self.run_transducer_loss(  scalar_t=torch.float16,
+                                                        fuse_softmax_backward=True,
+                                                        packed_input=False,
+                                                        for_vector_kernel=False)
+        self.assertTrue(torch.allclose(self.loss_ref, loss_tst, atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(self.grad_ref, grad_tst, atol=1e-4, rtol=1e-3))
+
+    def test_transducer_loss_fp16_backward_fusion_packed(self):
+        loss_tst, grad_tst = self.run_transducer_loss(  scalar_t=torch.float16,
+                                                        fuse_softmax_backward=True,
+                                                        packed_input=True,
+                                                        for_vector_kernel=False)
+        self.assertTrue(torch.allclose(self.loss_ref, loss_tst, atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(self.grad_ref, grad_tst, atol=1e-4, rtol=1e-3))
+
+    def test_transducer_loss_fp16_backward_fusion_packed_vec(self):
+        loss_tst, grad_tst = self.run_transducer_loss(  scalar_t=torch.float16,
+                                                        fuse_softmax_backward=True,
+                                                        packed_input=True,
+                                                        for_vector_kernel=True)
+        self.assertTrue(torch.allclose(self.loss_ref, loss_tst, atol=1e-5, rtol=1e-5))
+        self.assertTrue(torch.allclose(self.grad_ref, grad_tst, atol=1e-4, rtol=1e-3))
+
+
+
+if __name__ == '__main__':
+    unittest.main()
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/transducer_ref.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/transducer_ref.py
new file mode 100644
index 0000000000..75e65d4951
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/test/transducer/transducer_ref.py
@@ -0,0 +1,112 @@
+import torch
+import numpy as np
+import pdb
+
+def transducer_loss_reference(x, label, f_len, y_len, blank_idx, loss_grad):
+    def log_sum_exp(a, b):
+        if (a >= b):
+            return a + torch.log(1 + torch.exp(b-a))
+        else:
+            return b + torch.log(1 + torch.exp(a-b))
+
+    def forward_alpha(x, label, f_len, y_len, blank_idx):
+        B, T, U, V = x.size()
+        acc_t = torch.float32 if x.dtype in [torch.float16, torch.float32] else x.dtype
+        alpha = torch.zeros((B, T, U), dtype=acc_t, device=x.device)
+        for b in range(B):
+            alpha[b, 0, 0] = 0
+            for t in range(1, f_len[b]):
+                alpha[b, t, 0] = alpha[b, t-1, 0] + x[b, t-1, 0, blank_idx]
+            for u in range(1, y_len[b]+1):
+                alpha[b, 0, u] = alpha[b, 0, u-1] + x[b, 0, u-1, label[b, u-1]]
+            for t in range(1, f_len[b]):
+                for u in range(1, y_len[b]+1):
+                    curr_ = alpha[b, t-1, u] + x[b, t-1, u, blank_idx]
+                    next_ = alpha[b, t, u-1] + x[b, t, u-1, label[b, u-1]]
+                    alpha[b, t, u] = log_sum_exp(curr_, next_) 
+        return alpha
+
+    def forward_beta(x, label, f_len, y_len, blank_idx):
+        B, T, U, V = x.shape
+        acc_t = torch.float32 if x.dtype in [torch.float16, torch.float32] else x.dtype
+        beta = torch.zeros((B, T, U), dtype=acc_t, device=x.device)
+        for b in range(B):
+            beta[b, f_len[b]-1, y_len[b]] = x[b, f_len[b]-1, y_len[b], blank_idx]
+            for t in range(f_len[b]-2, -1, -1):
+                beta[b, t, y_len[b]] = beta[b, t+1, y_len[b]] + x[b, t, y_len[b], blank_idx] 
+            for u in range(y_len[b]-1, -1, -1):
+                beta[b, f_len[b]-1, u] = beta[b, f_len[b]-1, u+1] + x[b, f_len[b]-1, u, label[b, u]]
+            for t in range(f_len[b]-2, -1, -1):
+                for u in range(y_len[b]-1, -1, -1):
+                    curr_ = beta[b, t+1, u] + x[b, t, u, blank_idx] 
+                    next_ = beta[b, t, u+1] + x[b, t, u, label[b, u]]
+                    beta[b, t, u] = log_sum_exp(curr_, next_) 
+        return beta
+
+    def backward(x, label, f_len, y_len, alpha, beta, loss_grad, blank_idx):
+        grad = torch.zeros_like(x)
+        B, T, U, V = x.size()
+        for b in range(B):
+            common_factor = torch.log(loss_grad[b]) + alpha - beta[b, 0, 0]
+            # next
+            for u in range(y_len[b]):
+                grad[b, :f_len[b], u, label[b, u]] = -torch.exp(common_factor[b, :f_len[b], u] 
+                                                        + beta[b, :f_len[b], u+1] 
+                                                        + x[b, :f_len[b], u, label[b, u]])
+
+            # current
+            grad[b, :f_len[b]-1, :y_len[b]+1, blank_idx] \
+                = -torch.exp(common_factor[b, :f_len[b]-1, :y_len[b]+1] 
+                    + beta[b, 1:f_len[b], :y_len[b]+1] 
+                    + x[b, :f_len[b]-1, :y_len[b]+1, blank_idx])
+
+            grad[b, f_len[b]-1, y_len[b], blank_idx] = -torch.exp(common_factor[b, f_len[b]-1, y_len[b]]
+                                                         + x[b, f_len[b]-1, y_len[b], blank_idx])
+     
+        return grad
+
+    x_log = torch.nn.functional.log_softmax(x, dim=-1)
+    alpha = forward_alpha(x_log, label, f_len, y_len, blank_idx)
+    beta = forward_beta(x_log, label, f_len, y_len, blank_idx)
+    grad = backward(x_log, label, f_len, y_len, alpha, beta, 
+                        loss_grad, blank_idx)
+    x_log.backward(grad)
+    loss = -beta[:, 0, 0]
+    loss = loss.to(x.dtype)
+    return alpha, beta, x.grad, loss
+
+
+def transducer_joint_reference(f, g, h_grad, f_len, g_len, pack_output, relu, dropout, 
+                                dropout_prob=0, mask=None):
+    if dropout and mask == None:
+        raise NotImplementedError("mask needs to supplied to test dropout.")
+    B, T, H = f.size()
+    U = g.size(1)
+    f_expand = f.unsqueeze(dim=2)
+    g_expand = g.unsqueeze(dim=1)
+    h = f_expand + g_expand
+    if relu:
+        h = torch.nn.functional.relu(h)
+    if dropout:
+        h *= mask
+        scale = 1/(1-dropout_prob)
+        h *= scale
+    h.backward(h_grad)
+
+    if pack_output == False:
+        # intentionally set don't-care region to -1 to test if transducer joint
+        # write these regions to avoid NaN and inf
+        for b in range(B):
+            h[b, f_len[b]:] = -1
+            h[b, :, g_len[b]:] = -1
+
+        return h, f.grad, g.grad 
+
+    # packing
+    list_to_pack = []
+    for b in range(B):
+        list_to_pack.append(h[b, :f_len[b], :g_len[b], :].reshape(-1, H))
+    h_packed = torch.cat(list_to_pack)
+    return h_packed, f.grad, g.grad
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/transducer/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/transducer/__init__.py
new file mode 100644
index 0000000000..2c294687cf
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/transducer/__init__.py
@@ -0,0 +1,2 @@
+from .transducer import TransducerJoint
+from .transducer import TransducerLoss
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/transducer/transducer.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/transducer/transducer.py
new file mode 100644
index 0000000000..11ee457f40
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/transducer/transducer.py
@@ -0,0 +1,195 @@
+import torch
+import transducer_loss_cuda
+import transducer_joint_cuda
+
+class TransducerJoint(torch.nn.Module):
+    """Transducer joint
+    Detail of this loss function can be found in: Sequence Transduction with Recurrent Neural 
+    Networks
+
+    Arguments:
+        pack_output (bool, optional): whether to pack the output in a compact form with don't-care 
+        data being removed. (default: False)
+        relu (bool, optional): apply ReLU to the output of the joint operation. Requires opt=1  
+        (default: False)
+        dropout (bool, optional): apply dropout to the output of the joint operation. Requires opt=1  
+        (default: False)
+        opt (int, optional): pick the optimization level in [0, 1]. opt=1 picks a tiled algorithm. 
+            (default: 1)
+        fwd_tile_size (int, optional): tile size used in forward operation. This argument will be 
+        ignored if opt != 1. (default: 4) 
+        dropout_prob (float, optional): dropout probability. (default: 0.0)
+        probe_mask (bool, optional): a flag used to probe the mask generated by ReLU and/or dropout
+        operation. When this argument is set to True, the mask can be accessed through 
+        self.mask_probe. (default: false)
+    """
+
+    def __init__(self, pack_output=False, relu=False, dropout=False, opt=1, fwd_tile_size=4, 
+                    dropout_prob=0, probe_mask=False):
+        super(TransducerJoint, self).__init__() 
+        self.pack_output = pack_output
+        self.relu = relu
+        self.dropout = dropout
+        self.dropout_prob = dropout_prob
+        self.opt = opt
+        self.fwd_tile_size = fwd_tile_size
+        self.dummy_batch_offset = torch.empty(0)
+        masked = self.relu or self.dropout
+        self.mask_probe = [] if masked and probe_mask else None
+        if masked and opt != 1:
+            raise NotImplementedError("ReLU and dropout fusion is only supported with opt=1")
+
+
+    def forward(self, f, g, f_len, g_len, batch_offset=None, packed_batch=0):
+        """Forward operation of transducer joint
+
+        Arguments:
+            f (tensor): transcription vector from encode block of shape (B, T, H).
+            g (tensor): prediction vector form predict block of shape (B, U, H).
+            f_len (tensor): length of transcription vector for each batch.
+            g_len (tensor): length of prediction vector minus 1 for each batch.
+            batch_offset (tensor, optional): tensor containing the offset of each batch
+                in the results. For example, batch offset can be obtained from: 
+                batch_offset = torch.cumsum(f_len*g_len, dim=0)
+                This argument is required if pack_output == True, and is ignored if 
+                pack_output == False. (default: None)
+            packed_batch (int, optional): the batch size after packing. This argument is 
+                ignored if pack_output == False. (default: 0)
+        """
+        my_batch_offset = batch_offset if self.pack_output else self.dummy_batch_offset
+        if self.pack_output and (batch_offset is None or packed_batch == 0):
+            raise Exception("Please specify batch_offset and packed_batch when packing is enabled")
+        dropout =  self.dropout and self.training    # only dropout for training
+        return TransducerJointFunc.apply(f, g, f_len, g_len, self.pack_output, self.relu, dropout, 
+                                            my_batch_offset, packed_batch, self.opt, 
+                                            self.fwd_tile_size, self.dropout_prob, self.mask_probe)
+
+
+class TransducerLoss(torch.nn.Module):
+    """Transducer loss
+    Detail of this loss function can be found in: Sequence Transduction with Recurrent Neural 
+    Networks
+
+    Arguments:
+        fuse_softmax_backward (bool, optional) whether to fuse the backward of transducer loss with
+            softmax. (default: True)
+        opt (int, optional): pick the optimization level in [0, 1]. opt=1 picks a more optimized 
+            algorithm. In some cases, opt=1 might fall back to opt=0. (default: 1)
+        packed_input (bool, optional): whether to pack the output in a compact form with don't-care 
+        data being removed. (default: False)
+    """
+    def __init__(self, fuse_softmax_backward=True, opt=1, packed_input=False):
+        super(TransducerLoss, self).__init__() 
+        self.fuse_softmax_backward = fuse_softmax_backward
+        self.opt = opt
+        self.packed_input = packed_input
+        self.dummy_batch_offset = torch.empty(0)
+
+
+    def forward(self, x, label, f_len, y_len, blank_idx, batch_offset=None, max_f_len=None, 
+                debug_list=None):
+        """Forward operation of transducer joint
+
+        Arguments:
+            x (tensor): input tensor to the loss function with a shape of (B, T, U, H).
+            label (tensor): labels for the input data.
+            f_len (tensor): lengths of the inputs in the time dimension for each batch.
+            y_len (tensor): lengths of the labels for each batch.
+            blank_idx (int): index for the null symbol.
+            batch_offset (tensor, optional): tensor containing the offset of each batch
+                in the input. For example, batch offset can be obtained from: 
+                batch_offset = torch.cumsum(f_len*(y_len+1), dim=0)
+                This argument is required if packed_input == True, and is ignored if 
+                packed_input == False. (default: None)
+            max_f_len (int, optional): maximum length of the input in the time dimension.
+                For example, it can be obtained as 
+                max_f_len = max(f_len)
+                This argument is required if packed_input == True, and is ignored if 
+                packed_input == False. (default: None)
+                (default: None)
+            debug_list (list, optional): when an empty list is supplied, Alpha and Beta generated 
+                in the forward operation will be attached to this list for debug purpose. 
+                (default: None)
+        """
+        if self.packed_input:
+            if batch_offset is None or max_f_len is None:
+                raise Exception("Please specify batch_offset and max_f_len when packing is \
+                                    enabled") 
+            my_batch_offset = batch_offset
+            my_max_f_len = max_f_len
+        else:
+            my_batch_offset = self.dummy_batch_offset
+            my_max_f_len = x.size(1)
+        return TransducerLossFunc.apply(x, label, f_len, y_len, my_batch_offset, my_max_f_len, 
+                                            blank_idx, self.fuse_softmax_backward, debug_list, 
+                                            self.opt, self.packed_input)
+
+class TransducerLossFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, x, label, f_len, y_len, batch_offset, max_f_len, blank_idx, 
+                fuse_softmax_backward, debug_list, opt, packed_input):
+        if fuse_softmax_backward == False:
+            with torch.enable_grad():
+                x = torch.nn.functional.log_softmax(x, dim=-1)
+        else:
+            x = torch.nn.functional.log_softmax(x, dim=-1)
+        alpha, beta, loss = transducer_loss_cuda.forward(   x, label, f_len, y_len, batch_offset, 
+                                                            max_f_len, blank_idx, opt, packed_input)
+        if debug_list == []:
+            debug_list += [alpha, beta]
+        ctx.save_for_backward(x, alpha, beta, f_len, y_len, label, batch_offset)
+        ctx.blank_idx = blank_idx
+        ctx.fuse_softmax_backward = fuse_softmax_backward
+        ctx.opt = opt
+        ctx.packed_input = packed_input
+        ctx.max_f_len = max_f_len
+        return loss
+
+    @staticmethod
+    def backward(ctx, loss_grad):
+        x, alpha, beta, f_len, y_len, label, batch_offset = ctx.saved_tensors
+        x_grad = transducer_loss_cuda.backward( x, loss_grad, alpha, beta, f_len, y_len, label, 
+                                                batch_offset, ctx.max_f_len, ctx.blank_idx, ctx.opt, 
+                                                ctx.fuse_softmax_backward, ctx.packed_input)
+        if ctx.fuse_softmax_backward == False:
+            x_grad = x.backward(x_grad)
+        return x_grad, None, None, None, None, None, None, None, None, None, None
+
+class TransducerJointFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, f, g, f_len, g_len, pack_output, relu, dropout, batch_offset, packed_batch, 
+                opt, fwd_tile_size, dropout_prob, mask_probe):
+        h = transducer_joint_cuda.forward(f, g, f_len, g_len, batch_offset, packed_batch, opt, 
+                                            pack_output, relu, dropout, dropout_prob, fwd_tile_size)
+        masked = relu or dropout
+        if masked:
+            ctx.save_for_backward(h[1], f_len, g_len, batch_offset)
+            if mask_probe is not None:
+                mask_probe.append(h[1])
+        else:
+            ctx.save_for_backward(f_len, g_len, batch_offset)
+
+        ctx.pack_output = pack_output
+        ctx.masked = relu or dropout
+        ctx.max_f_len = f.size(1)
+        ctx.max_g_len = g.size(1)
+        ctx.scale = 1 / (1-dropout_prob) if dropout and dropout_prob != 1 else 1
+        return h[0]
+
+    @staticmethod
+    def backward(ctx, loss_grad):
+        if ctx.masked:
+            mask, f_len, g_len, batch_offset = ctx.saved_tensors
+            inp = [loss_grad, mask]
+        else:
+            f_len, g_len, batch_offset = ctx.saved_tensors
+            inp = [loss_grad]
+
+        f_grad, g_grad = transducer_joint_cuda.backward(    inp, f_len, g_len, batch_offset, 
+                                                            ctx.max_f_len, ctx.max_g_len, 
+                                                            ctx.pack_output, ctx.scale)
+
+        return f_grad, g_grad, None, None, None, None, None, None, None, None, None, None, None, \
+                None, None, None
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/xentropy/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/xentropy/__init__.py
new file mode 100644
index 0000000000..7dff6a27a5
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/xentropy/__init__.py
@@ -0,0 +1,9 @@
+try:
+    import torch
+    import xentropy_cuda
+    from .softmax_xentropy import SoftmaxCrossEntropyLoss
+    del torch
+    del xentropy_cuda
+    del softmax_xentropy
+except ImportError as err:
+    print("apex was installed without --xentropy flag, contrib.xentropy is not available")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/xentropy/softmax_xentropy.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/xentropy/softmax_xentropy.py
new file mode 100644
index 0000000000..33fbf8b216
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/contrib/xentropy/softmax_xentropy.py
@@ -0,0 +1,28 @@
+import torch
+import xentropy_cuda
+
+class SoftmaxCrossEntropyLoss(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, logits, labels, smoothing=0.0, padding_idx=0, half_to_float=False):
+        losses, max_log_sum_exp = xentropy_cuda.forward(
+            logits, labels, smoothing, half_to_float)
+        losses.masked_fill_(labels==padding_idx, 0)
+
+        ctx.save_for_backward(logits, max_log_sum_exp, labels,
+            torch.FloatTensor([smoothing]),
+            torch.LongTensor([padding_idx]))
+
+        return losses
+
+    @staticmethod
+    def backward(ctx, grad_loss):
+        logits, max_log_sum_exp, labels, smoothing, padding_idx = ctx.saved_tensors
+
+        if not grad_loss.is_contiguous():
+            grad_loss = grad_loss.contiguous()
+        grad_loss.masked_fill_(labels==padding_idx.item(), 0)
+        grad_logits = xentropy_cuda.backward(
+            grad_loss.contiguous(), logits, max_log_sum_exp,
+            labels, smoothing.item())
+
+        return grad_logits, None, None, None, None
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/README.md
new file mode 100644
index 0000000000..941de17947
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/README.md
@@ -0,0 +1,16 @@
+fp16_optimizer.py contains `FP16_Optimizer`, a Python class designed to wrap an existing Pytorch optimizer and automatically enable master parameters and loss scaling in a manner transparent to the user.  To use `FP16_Optimizer`, only two lines of one's Python model need to change.
+
+#### [FP16_Optimizer API documentation](https://nvidia.github.io/apex/fp16_utils.html#automatic-management-of-master-params-loss-scaling)
+
+#### [Simple examples with FP16_Optimizer](https://github.com/NVIDIA/apex/tree/master/examples/FP16_Optimizer_simple)
+
+#### [Imagenet with FP16_Optimizer](https://github.com/NVIDIA/apex/tree/master/examples/imagenet)
+
+#### [word_language_model with FP16_Optimizer](https://github.com/NVIDIA/apex/tree/master/examples/word_language_model)
+
+
+fp16_util.py contains a number of utilities to manually manage master parameters and loss scaling, if the user chooses.  
+
+#### [Manual management documentation](https://nvidia.github.io/apex/fp16_utils.html#manual-master-parameter-management)
+
+The [Imagenet with FP16_Optimizer](https://github.com/NVIDIA/apex/tree/master/examples/imagenet) and [word_language_model with FP16_Optimizer](https://github.com/NVIDIA/apex/tree/master/examples/word_language_model) directories also contain `main.py` files that demonstrate manual management of master parameters and static loss scaling.  These examples illustrate what sort of operations `FP16_Optimizer` is performing automatically.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/__init__.py
new file mode 100644
index 0000000000..c7bb1f537e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/__init__.py
@@ -0,0 +1,16 @@
+from .fp16util import (
+    BN_convert_float,
+    network_to_half,
+    prep_param_lists,
+    model_grads_to_master_grads,
+    master_params_to_model_params,
+    tofp16,
+    to_python_float,
+    clip_grad_norm,
+    convert_module,
+    convert_network,
+    FP16Model,
+)
+
+from .fp16_optimizer import FP16_Optimizer
+from .loss_scaler import LossScaler, DynamicLossScaler
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/fp16_optimizer.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/fp16_optimizer.py
new file mode 100644
index 0000000000..4f830fcb3f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/fp16_optimizer.py
@@ -0,0 +1,554 @@
+import torch
+from torch import nn
+from torch.autograd import Variable
+from torch.nn.parameter import Parameter
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+from ..amp._amp_state import _amp_state, maybe_print
+from ..amp.scaler import LossScaler
+from ..multi_tensor_apply import multi_tensor_applier
+from .fp16util import model_grads_to_master_grads, master_params_to_model_params, clip_grad_norm
+
+# TODO:  Update overflow check + downscale to use Carl's fused kernel.
+class FP16_Optimizer(object):
+    def __init__(self, 
+                 init_optimizer, 
+                 static_loss_scale=1.0, 
+                 dynamic_loss_scale=False,
+                 dynamic_loss_args=None,
+                 verbose=True):
+        print("Warning:  FP16_Optimizer is deprecated and dangerous, and will be deleted soon.  "
+              "If it still works, you're probably getting lucky.  "
+              "For mixed precision, use the documented API https://nvidia.github.io/apex/amp.html, with opt_level=O1.")
+
+        if not torch.cuda.is_available:
+            raise SystemError("Cannot use fp16 without CUDA.")
+
+        self.verbose = verbose
+
+        self.optimizer = init_optimizer
+        # init_state_dict sets up an alternative way to cast per-param state tensors.
+        # Stashing here in case https://github.com/pytorch/pytorch/issues/7733 makes it necessary.
+        # init_state_dict = init_optimizer.state_dict()
+
+        self.fp16_groups = []
+        self.fp32_from_fp16_groups = []
+        self.fp32_from_fp32_groups = []
+        for i, param_group in enumerate(self.optimizer.param_groups):
+            self.maybe_print("FP16_Optimizer processing param group {}:".format(i))
+            fp16_params_this_group = []
+            fp32_params_this_group = []
+            fp32_from_fp16_params_this_group = []
+            for i, param in enumerate(param_group['params']):
+                if param.requires_grad:
+                    if param.type() == 'torch.cuda.HalfTensor':
+                        self.maybe_print("FP16_Optimizer received torch.cuda.HalfTensor with {}"
+                                         .format(param.size()))
+                        fp16_params_this_group.append(param)
+                        master_param = param.detach().clone().float()
+                        master_param.requires_grad = True
+                        param_group['params'][i] = master_param
+                        fp32_from_fp16_params_this_group.append(master_param)
+                        # Reset existing state dict key to the new master param.
+                        # We still need to recast per-param state tensors, if any, to FP32.
+                        if param in self.optimizer.state:
+                           self.optimizer.state[master_param] = self.optimizer.state.pop(param) 
+                    elif param.type() == 'torch.cuda.FloatTensor':
+                        self.maybe_print("FP16_Optimizer received torch.cuda.FloatTensor with {}"
+                                         .format(param.size()))
+                        fp32_params_this_group.append(param)
+                        param_group['params'][i] = param
+                    else:
+                        raise TypeError("Wrapped parameters must be either "
+                                        "torch.cuda.FloatTensor or torch.cuda.HalfTensor. "  
+                                        "Received {}".format(param.type()))
+            
+            self.fp16_groups.append(fp16_params_this_group)
+            self.fp32_from_fp16_groups.append(fp32_from_fp16_params_this_group)
+            self.fp32_from_fp32_groups.append(fp32_params_this_group)
+
+        self.all_fp16_params = []
+        for group in self.fp16_groups:
+            self.all_fp16_params += group
+
+        self.all_fp32_from_fp16_params = []
+        for group in self.fp32_from_fp16_groups:
+            self.all_fp32_from_fp16_params += group
+
+        self.all_fp32_from_fp32_params = []
+        for group in self.fp32_from_fp32_groups:
+            self.all_fp32_from_fp32_params += group
+
+        # Leverage state_dict() and load_state_dict() to recast preexisting per-param state tensors
+        self.optimizer.load_state_dict(self.optimizer.state_dict())
+        # alternative way to cast per-param state tensors:
+        # self.optimizer.load_state_dict(init_state_dict)
+
+        if dynamic_loss_scale:
+            self.dynamic_loss_scale = True
+            if dynamic_loss_args is not None:
+                self.loss_scaler = LossScaler("dynamic", **dynamic_loss_args)
+            else:
+                self.loss_scaler = LossScaler("dynamic")
+        else:
+            self.dynamic_loss_scale = False
+            self.loss_scaler = LossScaler(static_loss_scale)
+
+        self.overflow = False
+        self.first_closure_call_this_step = True
+
+        self.clip_grad_norm = clip_grad_norm
+
+        # TODO:  Centralize exposure and import error checking for the C backend.
+        if multi_tensor_applier.available:
+            import amp_C
+            self.multi_tensor_scale = amp_C.multi_tensor_scale
+            self._dummy_overflow_buf = torch.cuda.IntTensor([0]);
+
+    # Having self.maybe_print distinct from _amp_state.maybe_print is another artifact
+    # of having to support FP16_Optimizer separately, for the time being.
+    def maybe_print(self, msg):
+        if self.verbose:
+            print(msg)
+            
+    def __getstate__(self):
+        raise RuntimeError("FP16_Optimizer should be serialized using state_dict().")
+
+    def __setstate__(self, state):
+        raise RuntimeError("FP16_Optimizer should be deserialized using load_state_dict().")
+
+    def zero_grad(self, set_grads_to_None=False):
+        """
+        Zero fp32 and fp16 parameter grads.
+        """
+        # In principle, only the .grad attributes of the model params need to be zeroed,
+        # because gradients are copied into the FP32 master params.  However, we zero
+        # all gradients owned by the optimizer, just to be safe:
+        for group in self.optimizer.param_groups:
+             for p in group['params']:
+                 if set_grads_to_None:
+                     p.grad = None
+                 else:
+                     if p.grad is not None:
+                         p.grad.detach_()
+                         p.grad.zero_()
+
+        # Zero fp16 gradients owned by the model:
+        for fp16_group in self.fp16_groups:
+            for param in fp16_group:
+                if set_grads_to_None:
+                    param.grad = None
+                else:
+                    if param.grad is not None:
+                        param.grad.detach_() # as in torch.optim.optimizer.zero_grad()
+                        param.grad.zero_()
+
+    # Should not be used anymore.
+    # def _check_overflow(self):
+    #     params = []
+    #     for group in self.fp16_groups:
+    #         for param in group:
+    #             params.append(param)
+    #     for group in self.fp32_from_fp32_groups:
+    #         for param in group:
+    #             params.append(param)
+    #     self.overflow = self.loss_scaler.has_overflow(params)
+
+    # def _update_scale(self, has_overflow=False):
+    #     self.loss_scaler.update_scale(has_overflow)
+
+    def _master_params_to_model_params(self):
+        if multi_tensor_applier.available:
+            if len(self.all_fp16_params) > 0:
+                multi_tensor_applier(
+                    self.multi_tensor_scale,
+                    self._dummy_overflow_buf,
+                    [self.all_fp32_from_fp16_params, self.all_fp16_params],
+                    1.0)
+        else:
+            for fp16_group, fp32_from_fp16_group in zip(self.fp16_groups, self.fp32_from_fp16_groups):
+                master_params_to_model_params(fp16_group, fp32_from_fp16_group)
+
+    # To consider:  Integrate distributed with this wrapper by registering a hook on each variable
+    # that does the overflow check, gradient copy + downscale, and fp32 allreduce in a different stream.
+    # def _model_grads_to_master_grads(self):
+    #     for fp16_group, fp32_from_fp16_group in zip(self.fp16_groups, self.fp32_from_fp16_groups):
+    #         model_grads_to_master_grads(fp16_group, fp32_from_fp16_group)
+
+    # def _downscale_master(self):
+    #     if self.loss_scale != 1.0:
+    #         for group in self.optimizer.param_groups:
+    #             for param in group['params']:
+    #                 if param.grad is not None:
+    #                     param.grad.data.mul_(1./self.loss_scale)
+
+    def clip_master_grads(self, max_norm, norm_type=2):
+        """
+        Clips fp32 master gradients via ``torch.nn.utils.clip_grad_norm``.
+
+        Args:
+            max_norm (float or int): max norm of the gradients
+            norm_type (float or int): type of the used p-norm. Can be ``'inf'`` for
+                infinity norm.
+
+        Returns:
+            Total norm of the current fp32 gradients (viewed as a single vector).
+
+        .. warning::
+            Returns -1 if the most recently computed fp16 gradients overflowed (that is, if ``self.overflow`` is ``True``).
+        """
+        if not self.overflow:
+            fp32_params = []
+            for param_group in self.optimizer.param_groups:
+                for param in param_group['params']:
+                    fp32_params.append(param)
+            return self.clip_grad_norm(fp32_params, max_norm, norm_type)
+        else:
+            return -1
+
+    def state_dict(self):
+        """
+        Returns a dict containing the current state of this :class:`FP16_Optimizer` instance.
+        This dict contains attributes of :class:`FP16_Optimizer`, as well as the state_dict
+        of the contained Pytorch optimizer.
+        Example::
+
+            checkpoint = {}
+            checkpoint['model'] = model.state_dict()
+            checkpoint['optimizer'] = optimizer.state_dict()
+            torch.save(checkpoint, "saved.pth")
+        """
+        state_dict = {}
+        state_dict['loss_scaler'] = self.loss_scaler
+        state_dict['dynamic_loss_scale'] = self.dynamic_loss_scale
+        state_dict['overflow'] = self.overflow
+        state_dict['first_closure_call_this_step'] = self.first_closure_call_this_step
+        state_dict['optimizer_state_dict'] = self.optimizer.state_dict()
+        state_dict['fp32_from_fp16'] = self.fp32_from_fp16_groups
+        return state_dict
+
+    def load_state_dict(self, state_dict):
+        """
+        Loads a state_dict created by an earlier call to state_dict(). 
+        If ``fp16_optimizer_instance`` was constructed from some ``init_optimizer``, 
+        whose parameters in turn came from ``model``, it is expected that the user 
+        will call ``model.load_state_dict()`` before
+        ``fp16_optimizer_instance.load_state_dict()`` is called.
+
+        Example::
+
+            model = torch.nn.Linear(D_in, D_out).cuda().half()
+            optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
+            optimizer = FP16_Optimizer(optimizer, static_loss_scale = 128.0)
+            ...
+            checkpoint = torch.load("saved.pth")
+            model.load_state_dict(checkpoint['model'])
+            optimizer.load_state_dict(checkpoint['optimizer'])
+        """
+        # I think it should actually be ok to reload the optimizer before the model.
+        self.loss_scaler = state_dict['loss_scaler']
+        self.dynamic_loss_scale = state_dict['dynamic_loss_scale']
+        self.overflow = state_dict['overflow']
+        self.first_closure_call_this_step = state_dict['first_closure_call_this_step']
+        self.optimizer.load_state_dict(state_dict['optimizer_state_dict'])
+        # At this point, the optimizer's references to the model's fp32 parameters are up to date.
+        # The optimizer's hyperparameters and internal buffers are also up to date.  
+        # However, the fp32 master copies of the model's fp16 params stored by the optimizer are still
+        # out of date.  There are two options.  
+        # 1:  Refresh the master params from the model's fp16 params.  
+        # This requires less storage but incurs precision loss.
+        # 2:  Save and restore the fp32 master copies separately.
+        # We choose option 2.
+        # 
+        # Pytorch Optimizer.load_state_dict casts saved buffers (e.g. momentum) to the type and device 
+        # of their associated parameters, because it's possible those buffers might not exist yet in 
+        # the current optimizer instance.  In our case, as long as the current FP16_Optimizer has been 
+        # constructed in the same way as the one whose state_dict we are loading, the same master params
+        # are guaranteed to exist, so we can just copy_() from the saved master params.
+        for current_group, saved_group in zip(self.fp32_from_fp16_groups, state_dict['fp32_from_fp16']):
+            for current, saved in zip(current_group, saved_group):
+                current.data.copy_(saved.data)
+
+    def step(self, closure=None): # could add clip option.
+        """
+        If no closure is supplied, :attr:`step` should be called after 
+        ``fp16_optimizer_obj.backward(loss)``.
+        :attr:`step` updates the fp32 master copy of parameters using the optimizer supplied to
+        :class:`FP16_Optimizer`'s constructor, then copies the updated fp32 params into the fp16 params
+        originally referenced by :class:`FP16_Optimizer`'s constructor, so the user may immediately run
+        another forward pass using their model.
+
+        If a closure is supplied, :attr:`step` may be called without a prior call to 
+        :attr:`backward(loss)`.
+        This control flow is identical to `ordinary Pytorch optimizer use`_ with closures.
+        However, the user should take care that any ``loss.backward()`` call within the closure
+        has been replaced by ``fp16_optimizer_obj.backward(loss)``.
+
+        Args:
+           closure (optional):  Closure that will be supplied to the underlying optimizer originally passed to :class:`FP16_Optimizer`'s constructor.  closure should call :attr:`zero_grad()` on the :class:`FP16_Optimizer` object, compute the loss, call :attr:`backward(loss)`, and return the loss.
+
+        Example with closure::
+
+            # optimizer is assumed to be an FP16_Optimizer object, previously constructed from an 
+            # existing pytorch optimizer.
+            for input, target in dataset:
+                def closure():
+                    optimizer.zero_grad()
+                    output = model(input)
+                    loss = loss_fn(output, target)
+                    # loss.backward() becomes:
+                    optimizer.backward(loss)
+                    return loss
+                optimizer.step(closure)
+
+        .. warning::
+            Currently, calling :attr:`step` with a closure is not compatible with dynamic loss scaling.
+
+        .. _`ordinary Pytorch optimizer use`:
+            http://pytorch.org/docs/master/optim.html#optimizer-step-closure
+        """
+
+        scale = self.loss_scaler.loss_scale()
+        # To consider:  Should this be in step(), or update_master_grads?  It works either way,
+        # but I should make it consistent with the Amp control flow, which updates the scale
+        # during backward context manager exit.
+        # self._update_scale(self.overflow)
+
+        if self.overflow:
+            # Using _amp_state.maybe_print instead of self.print here is intentional.
+            maybe_print("Gradient overflow.  Skipping step, reducing " +
+                "loss scale to {}".format(self.loss_scaler.loss_scale()))
+            return
+        
+        if closure is not None:
+            retval = self._step_with_closure(closure)
+        else:
+            # torch.cuda.nvtx.range_push("pytorch optimizer step")
+            retval = self.optimizer.step()
+            # torch.cuda.nvtx.range_pop()
+
+        self._master_params_to_model_params()
+
+        return retval
+
+    def _step_with_closure(self, closure):
+        def wrapped_closure():
+            # helpful for debugging
+            # print("Calling wrapped_closure, first_closure_call_this_step = {}"
+            #       .format(self.first_closure_call_this_step))
+            if self.first_closure_call_this_step:
+                # We expect that the fp16 params are initially fresh on entering self.step(),
+                # so _master_params_to_model_params() is unnecessary the first time wrapped_closure()
+                # is called within self.optimizer.step().
+                self.first_closure_call_this_step = False
+            else:
+                # If self.optimizer.step() internally calls wrapped_closure more than once,
+                # it may update the fp32 params after each call.  However, self.optimizer 
+                # doesn't know about the fp16 params at all.  If the fp32 params get updated,
+                # we can't rely on self.optimizer to refresh the fp16 params.  We need
+                # to handle that manually:
+                self._master_params_to_model_params()
+            # Our API expects the user to give us ownership of the backward() call by
+            # replacing all calls to loss.backward() with optimizer.backward(loss).
+            # This requirement holds whether or not the call to backward() is made within a closure.
+            # If the user is properly calling optimizer.backward(loss) within "closure," 
+            # calling closure() here will give the fp32 master params fresh gradients
+            # for the optimizer to play with, so all wrapped_closure needs to do is call 
+            # closure() and return the loss.
+            temp_loss = closure() 
+            while(self.overflow):
+                scale = self.loss_scaler.loss_scale()
+                # self._update_scale(self.overflow) # now done at the end of backward
+                print("OVERFLOW within closure! Skipping step, reducing loss scale to {}".format(
+                      self.loss_scaler.loss_scale()))
+                temp_loss = closure()
+            return temp_loss
+
+        retval = self.optimizer.step(wrapped_closure)
+
+        self.first_closure_call_this_step = True
+
+        return retval
+
+    def backward(self, loss, update_master_grads=True, retain_graph=False):
+        """ 
+        :attr:`backward` performs the following conceptual steps:
+
+        1. fp32_loss = loss.float() (see first Note below)
+        2. scaled_loss = fp32_loss*loss_scale
+        3. scaled_loss.backward(), which accumulates scaled gradients into the ``.grad`` attributes of the model's leaves (which may be fp16, fp32, or a mixture, depending how your model was defined).
+        4. fp16 grads are then copied to the master params' ``.grad`` attributes (see second Note), which are guaranteed to be fp32.
+        5. Finally, master grads are divided by loss_scale.
+
+        In this way, after :attr:`backward`, the master params have fresh gradients,
+        and :attr:`step` may be called.
+
+        .. note::
+            :attr:`backward` internally converts the loss to fp32 before applying the loss scale.
+            This provides some additional safety against overflow if the user has supplied an 
+            fp16 loss value.  
+            However, for maximum overflow safety, the user should
+            compute the loss criterion (MSE, cross entropy, etc) in fp32 before supplying it to 
+            :attr:`backward`.
+
+        .. warning::
+            The gradients found in a model's leaves after the call to 
+            :attr:`backward` should not be regarded as valid in general, 
+            because it's possible 
+            they have been scaled (and in the case of dynamic loss scaling, 
+            the scale factor may change over time).  
+            If the user wants to inspect gradients after a call to :attr:`backward`,  
+            only the master gradients should be regarded as valid.  These can be retrieved via
+            :attr:`inspect_master_grad_data()`.
+
+        Args:
+            loss:  The loss output by the user's model.  loss may be either float or half (but see first Note above).
+            update_master_grads (bool, optional, default=True):  Option to copy fp16 grads to fp32 grads on this call.  By setting this to False, the user can delay the copy, which is useful to eliminate redundant fp16->fp32 grad copies if :attr:`backward` is being called on multiple losses in one iteration.  If set to False, the user becomes responsible for calling :attr:`update_master_grads` before calling :attr:`step`.
+            retain_graph (bool, optional, default=False):  Forwards the usual ``retain_graph=True`` option to the internal call to ``loss.backward``.  If ``retain_graph`` is being used to accumulate gradient values from multiple backward passes before calling ``optimizer.step``, passing ``update_master_grads=False`` is also recommended (see Example below).
+
+        Example::
+
+            # Ordinary operation:
+            optimizer.backward(loss)
+
+            # Naive operation with multiple losses (technically valid, but less efficient):
+            # fp32 grads will be correct after the second call,  but 
+            # the first call incurs an unnecessary fp16->fp32 grad copy.
+            optimizer.backward(loss1)
+            optimizer.backward(loss2)
+
+            # More efficient way to handle multiple losses:
+            # The fp16->fp32 grad copy is delayed until fp16 grads from all 
+            # losses have been accumulated.
+            optimizer.backward(loss1, update_master_grads=False)
+            optimizer.backward(loss2, update_master_grads=False)
+            optimizer.update_master_grads()
+        """ 
+        # To consider:  try multiple backward passes using retain_grad=True to find 
+        # a loss scale that works.  After you find a loss scale that works, do a final dummy
+        # backward pass with retain_graph=False to tear down the graph.  Doing this would avoid 
+        # discarding the iteration,  but probably wouldn't improve overall efficiency.  
+        scaled_loss = loss.float()*self.loss_scaler.loss_scale()
+        scaled_loss.backward(retain_graph=retain_graph)
+        if update_master_grads:
+            self.update_master_grads()
+
+    def update_master_grads(self):
+        # torch.cuda.nvtx.range_push("update_master_grads")
+        """
+        Copy the ``.grad`` attribute from stored references to fp16 parameters to 
+        the ``.grad`` attribute of the fp32 master parameters that are directly 
+        updated by the optimizer.  :attr:`update_master_grads` only needs to be called if
+        ``fp16_optimizer_obj.backward`` was called with ``update_master_grads=False``.
+        """
+        # if self.dynamic_loss_scale:
+        #     self._check_overflow()
+        #     if self.overflow: return
+        # self._model_grads_to_master_grads()
+        # self._downscale_master()
+        # Use the one-shot multi-tensor apply kernel
+        self.loss_scaler.clear_overflow_state()
+        if len(self.all_fp16_params) > 0:
+            # print("Model grads before")
+            # print([param.grad.data for param in self.all_fp16_params])
+            # I'm ONLY writing this as an incremental way to make some tests pass until
+            # I can refactor the tests as well.
+            # FP16_Optimizer should not be used by anyone.
+            model_grads = []
+            master_grads = []
+            for model_param, master_param in zip(self.all_fp16_params,
+                                                 self.all_fp32_from_fp16_params):
+                if model_param.grad is not None:
+                    model_grads.append(model_param.grad)
+                    if master_param.grad is None:
+                        master_param.grad = torch.empty_like(master_param)
+                    master_grads.append(master_param.grad)
+            self.loss_scaler.unscale(
+                model_grads,
+                master_grads,
+                self.loss_scaler.loss_scale())
+            # print("Master grads after")
+            # print([param.grad.data for param in self.all_fp32_from_fp16_params])
+        if len(self.all_fp32_from_fp32_params) > 0:
+            model_grads = []
+            master_grads = []
+            for model_param, master_param in zip(self.all_fp32_from_fp32_params,
+                                                 self.all_fp32_from_fp32_params):
+                if model_param.grad is not None:
+                    model_grads.append(model_param.grad)
+                    master_grads.append(master_param.grad)
+            # print("Model grads before")
+            # print([param.grad.data for param in self.all_fp32_from_fp32_params])
+            self.loss_scaler.unscale(
+                model_grads,
+                master_grads,
+                self.loss_scaler.loss_scale())
+            # print("Master grads after")
+            # print([param.grad.data for param in self.all_fp32_from_fp32_params])
+        # quit()
+        self.overflow = self.loss_scaler.update_scale()
+        # torch.cuda.nvtx.range_pop()
+
+
+    def inspect_master_grad_data(self):
+        """
+        When running with :class:`FP16_Optimizer`, 
+        ``.grad`` attributes of a model's fp16 leaves should not be
+        regarded as truthful, because they might be scaled.  
+        After a call to :attr:`fp16_optimizer_obj.backward(loss)`, if no overflow was encountered,
+        the fp32 master params' ``.grad``
+        attributes will contain valid gradients properly divided by the loss scale.  However, 
+        because :class:`FP16_Optimizer` flattens some parameters, accessing them may be 
+        nonintuitive.  :attr:`inspect_master_grad_data`
+        allows those gradients to be viewed with shapes corresponding to their associated model leaves.
+
+        Returns:
+            List of lists (one list for each parameter group).  The list for each parameter group
+            is a list of the ``.grad.data`` attributes of the fp32 master params belonging to that group.                 
+        """
+        if self.overflow:
+            print("Warning:  calling FP16_Optimizer.inspect_master_grad_data while in an overflow state.  "
+                  "Gradients are currently invalid (may be inf, nan, or stale).  Returning None.")
+            return None
+        else:
+            # The optimizer owns only references to master params.
+            master_grads_data = []
+            for param_group in self.optimizer.param_groups:
+                master_grads_this_group = []
+                for param in param_group['params']:
+                    if param.grad is not None:
+                        master_grads_this_group.append(param.grad.data)
+                    else:
+                        master_grads_this_group.append(None)
+                master_grads_data.append(master_grads_this_group)
+            return master_grads_data
+
+
+    # Promote loss scale so it can be retrieved or set via "fp16_optimizer_instance.loss_scale"
+    def _get_loss_scale(self):
+        return self.loss_scaler.loss_scale()
+
+    def _set_loss_scale(self, value):
+        self.loss_scaler._loss_scale = value
+
+    loss_scale = property(_get_loss_scale, _set_loss_scale)
+
+    # Promote state so it can be retrieved or set via "fp16_optimizer_instance.state"
+    def _get_state(self):
+        return self.optimizer.state
+
+    def _set_state(self, value):
+        self.optimizer.state = value
+
+    state = property(_get_state, _set_state)
+
+    # Promote param_groups so it can be retrieved or set via "fp16_optimizer_instance.param_groups"
+    # (for example, to adjust the learning rate)
+    def _get_param_groups(self):
+        return self.optimizer.param_groups
+
+    def _set_param_groups(self, value):
+        self.optimizer.param_groups = value
+
+    param_groups = property(_get_param_groups, _set_param_groups)
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/fp16util.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/fp16util.py
new file mode 100644
index 0000000000..dcdc3447a4
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/fp16util.py
@@ -0,0 +1,187 @@
+import torch
+import torch.nn as nn
+from torch.autograd import Variable
+from torch._utils import _flatten_dense_tensors, _unflatten_dense_tensors
+
+
+class tofp16(nn.Module):
+    """
+    Utility module that implements::
+
+        def forward(self, input):
+            return input.half()
+    """
+
+    def __init__(self):
+        super(tofp16, self).__init__()
+
+    def forward(self, input):
+        return input.half()
+
+
+def BN_convert_float(module):
+    """
+    Utility function for network_to_half().
+
+    Retained for legacy purposes.
+    """
+    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm) and module.affine is True:
+        module.float()
+    for child in module.children():
+        BN_convert_float(child)
+    return module
+
+
+def network_to_half(network):
+    """
+    Convert model to half precision in a batchnorm-safe way.
+
+    Retained for legacy purposes. It is recommended to use FP16Model.
+    """
+    return nn.Sequential(tofp16(), BN_convert_float(network.half()))
+
+
+def convert_module(module, dtype):
+    """
+    Converts a module's immediate parameters and buffers to dtype.
+    """
+    for param in module.parameters(recurse=False):
+        if param is not None:
+            if param.data.dtype.is_floating_point:
+                param.data = param.data.to(dtype=dtype)
+            if param._grad is not None and param._grad.data.dtype.is_floating_point:
+                param._grad.data = param._grad.data.to(dtype=dtype)
+
+    for buf in module.buffers(recurse=False):
+        if buf is not None and buf.data.dtype.is_floating_point:
+            buf.data = buf.data.to(dtype=dtype)
+
+
+def convert_network(network, dtype):
+    """
+    Converts a network's parameters and buffers to dtype.
+    """
+    for module in network.modules():
+        if isinstance(module, torch.nn.modules.batchnorm._BatchNorm) and module.affine is True:
+            continue
+        convert_module(module, dtype)
+        if isinstance(module, torch.nn.RNNBase) or isinstance(module, torch.nn.modules.rnn.RNNBase):
+            module.flatten_parameters()
+    return network
+
+
+class FP16Model(nn.Module):
+    """
+    Convert model to half precision in a batchnorm-safe way.
+    """
+
+    def __init__(self, network):
+        super(FP16Model, self).__init__()
+        self.network = convert_network(network, dtype=torch.half)
+
+    def forward(self, *inputs):
+        inputs = tuple(t.half() for t in inputs)
+        return self.network(*inputs)
+
+
+def backwards_debug_hook(grad):
+    raise RuntimeError("master_params recieved a gradient in the backward pass!")
+
+def prep_param_lists(model, flat_master=False):
+    """
+    Creates a list of FP32 master parameters for a given model, as in
+    `Training Neural Networks with Mixed Precision:  Real Examples`_.
+
+    Args:
+        model (torch.nn.Module): Existing Pytorch model
+        flat_master (bool, optional, default=False):  Flatten the master parameters into a single tensor, as a performance optimization.
+    Returns:
+        A tuple (``model_params``, ``master_params``). ``model_params`` is a list of the model's parameters for later use with :func:`model_grads_to_master_grads` and :func:`master_params_to_model_params`.  ``master_params`` is a list of FP32 master gradients.  If ``flat_master=True``, ``master_params`` will be a list with one element.
+
+    Example::
+
+        model_params, master_params = prep_param_lists(model)
+
+    .. warning::
+        Currently, if ``flat_master=True``, all the model's parameters must be the same type.  If the model has parameters of different types, use ``flat_master=False``, or use :class:`FP16_Optimizer`.
+
+    .. _`Training Neural Networks with Mixed Precision:  Real Examples`:
+        http://on-demand.gputechconf.com/gtc/2018/video/S81012/
+    """
+    model_params = [param for param in model.parameters() if param.requires_grad]
+
+    if flat_master:
+        # Give the user some more useful error messages
+        try:
+            # flatten_dense_tensors returns a contiguous flat array.
+            # http://pytorch.org/docs/master/_modules/torch/_utils.html
+            master_params = _flatten_dense_tensors([param.data for param in model_params]).float()
+        except:
+            print("Error in prep_param_lists:  model may contain a mixture of parameters "
+                      "of different types.  Use flat_master=False, or use F16_Optimizer.")
+            raise
+        master_params = torch.nn.Parameter(master_params)
+        master_params.requires_grad = True
+        # master_params.register_hook(backwards_debug_hook)
+        if master_params.grad is None:
+            master_params.grad = master_params.new(*master_params.size())
+        return model_params, [master_params]
+    else:
+        master_params = [param.clone().float().detach() for param in model_params]
+        for param in master_params:
+            param.requires_grad = True
+        return model_params, master_params
+
+
+def model_grads_to_master_grads(model_params, master_params, flat_master=False):
+    """
+    Copy model gradients to master gradients.  
+
+    Args:
+        model_params:  List of model parameters created by :func:`prep_param_lists`.
+        master_params:  List of FP32 master parameters created by :func:`prep_param_lists`.  If ``master_params`` was created with ``flat_master=True``, ``flat_master=True`` should also be supplied to :func:`model_grads_to_master_grads`.
+    """
+    if flat_master:
+        # The flattening may incur one more deep copy than is necessary.
+        master_params[0].grad.data.copy_(
+            _flatten_dense_tensors([p.grad.data for p in model_params]))
+    else:
+        for model, master in zip(model_params, master_params):
+            if model.grad is not None:
+                if master.grad is None:
+                    master.grad = Variable(master.data.new(*master.data.size()))
+                master.grad.data.copy_(model.grad.data)
+            else:
+                master.grad = None
+
+
+def master_params_to_model_params(model_params, master_params, flat_master=False):
+    """
+    Copy master parameters to model parameters.
+
+    Args:
+        model_params:  List of model parameters created by :func:`prep_param_lists`.
+        master_params:  List of FP32 master parameters created by :func:`prep_param_lists`.  If ``master_params`` was created with ``flat_master=True``, ``flat_master=True`` should also be supplied to :func:`master_params_to_model_params`.
+    """
+    if flat_master:
+        for model, master in zip(model_params, 
+                                 _unflatten_dense_tensors(master_params[0].data, model_params)):
+            model.data.copy_(master)
+    else:
+        for model, master in zip(model_params, master_params):
+            model.data.copy_(master.data)
+
+# Backward compatibility fixes
+
+def to_python_float(t):
+    if hasattr(t, 'item'):
+        return t.item()
+    else:
+        return t[0]
+
+TORCH_MAJOR = int(torch.__version__.split('.')[0])
+TORCH_MINOR = int(torch.__version__.split('.')[1])
+if TORCH_MAJOR == 0 and TORCH_MINOR <= 4:
+    clip_grad_norm = torch.nn.utils.clip_grad_norm
+else:
+    clip_grad_norm = torch.nn.utils.clip_grad_norm_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/loss_scaler.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/loss_scaler.py
new file mode 100644
index 0000000000..b9f32fe01f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fp16_utils/loss_scaler.py
@@ -0,0 +1,186 @@
+import torch
+
+# item() is a recent addition, so this helps with backward compatibility.
+def to_python_float(t):
+    if hasattr(t, 'item'):
+        return t.item()
+    else:
+        return t[0]
+
+class LossScaler:
+    """
+    Class that manages a static loss scale.  This class is intended to interact with
+    :class:`FP16_Optimizer`, and should not be directly manipulated by the user.
+
+    Use of :class:`LossScaler` is enabled via the ``static_loss_scale`` argument to 
+    :class:`FP16_Optimizer`'s constructor.
+
+    Args:
+        scale (float, optional, default=1.0):  The loss scale.
+    """
+
+    def __init__(self, scale=1):
+        self.cur_scale = scale
+
+    # `params` is a list / generator of torch.Variable
+    def has_overflow(self, params):
+        return False
+
+    # `x` is a torch.Tensor
+    def _has_inf_or_nan(x):
+        return False
+
+    def update_scale(self, overflow):
+        pass
+
+    @property
+    def loss_scale(self):
+        return self.cur_scale
+
+    def scale_gradient(self, module, grad_in, grad_out):
+        return tuple(self.loss_scale * g for g in grad_in)
+
+    def backward(self, loss, retain_graph=False):
+        scaled_loss = loss*self.loss_scale
+        scaled_loss.backward(retain_graph=retain_graph)
+
+class DynamicLossScaler:
+    """
+    Class that manages dynamic loss scaling.  It is recommended to use :class:`DynamicLossScaler`
+    indirectly, by supplying ``dynamic_loss_scale=True`` to the constructor of 
+    :class:`FP16_Optimizer`.  However, it's important to understand how :class:`DynamicLossScaler`
+    operates, because the default options can be changed using the
+    the ``dynamic_loss_args`` argument to :class:`FP16_Optimizer`'s constructor.
+
+    Loss scaling is designed to combat the problem of underflowing gradients encountered at long
+    times when training fp16 networks.  Dynamic loss scaling begins by attempting a very high loss
+    scale.  Ironically, this may result in OVERflowing gradients.  If overflowing gradients are
+    encountered, :class:`DynamicLossScaler` informs :class:`FP16_Optimizer` that an overflow has 
+    occurred.
+    :class:`FP16_Optimizer` then skips the update step for this particular iteration/minibatch,
+    and :class:`DynamicLossScaler` adjusts the loss scale to a lower value.  
+    If a certain number of iterations occur without overflowing gradients detected,
+    :class:`DynamicLossScaler` increases the loss scale once more.
+    In this way :class:`DynamicLossScaler` attempts to "ride the edge" of 
+    always using the highest loss scale possible without incurring overflow.
+
+    Args:
+        init_scale (float, optional, default=2**32):  Initial loss scale attempted by :class:`DynamicLossScaler.`
+        scale_factor (float, optional, default=2.0):  Factor used when adjusting the loss scale. If an overflow is encountered, the loss scale is readjusted to loss scale/``scale_factor``.  If ``scale_window`` consecutive iterations take place without an overflow, the loss scale is readjusted to loss_scale*``scale_factor``. 
+        scale_window (int, optional, default=1000):  Number of consecutive iterations without an overflow to wait before increasing the loss scale.
+    """
+
+    def __init__(self,
+                 init_scale=2**32,
+                 scale_factor=2.,
+                 scale_window=1000):
+        self.cur_scale = init_scale
+        self.cur_iter = 0
+        self.last_overflow_iter = -1
+        self.scale_factor = scale_factor
+        self.scale_window = scale_window
+
+    # `params` is a list / generator of torch.Variable
+    def has_overflow(self, params):
+        for p in params:
+            if p.grad is not None and DynamicLossScaler._has_inf_or_nan(p.grad.data):
+                return True
+
+        return False
+
+    # `x` is a torch.Tensor
+    def _has_inf_or_nan(x):
+        try:
+            # if x is half, the .float() incurs an additional deep copy, but it's necessary if 
+            # Pytorch's .sum() creates a one-element tensor of the same type as x 
+            # (which is true for some recent version of pytorch).
+            cpu_sum = float(x.float().sum())
+            # More efficient version that can be used if .sum() returns a Python scalar
+            # cpu_sum = float(x.sum())
+        except RuntimeError as instance:
+            # We want to check if inst is actually an overflow exception.
+            # RuntimeError could come from a different error.
+            # If so, we still want the exception to propagate.
+            if "value cannot be converted" not in instance.args[0]:
+                raise
+            return True
+        else:
+            if cpu_sum == float('inf') or cpu_sum == -float('inf') or cpu_sum != cpu_sum:
+                return True
+            return False
+
+    # `overflow` is boolean indicating whether the gradient overflowed
+    def update_scale(self, overflow):
+        if overflow:
+            # self.cur_scale /= self.scale_factor
+            self.cur_scale = max(self.cur_scale/self.scale_factor, 1)
+            self.last_overflow_iter = self.cur_iter
+        else:
+            if (self.cur_iter - self.last_overflow_iter) % self.scale_window == 0:
+                self.cur_scale *= self.scale_factor
+        self.cur_iter += 1
+
+    @property
+    def loss_scale(self):
+        return self.cur_scale
+
+    def scale_gradient(self, module, grad_in, grad_out):
+        return tuple(self.loss_scale * g for g in grad_in)
+
+    def backward(self, loss, retain_graph=False):
+        scaled_loss = loss*self.loss_scale
+        scaled_loss.backward(retain_graph=retain_graph)
+        
+##############################################################        
+# Example usage below here -- assuming it's in a separate file
+##############################################################
+"""
+TO-DO separate out into an example.
+if __name__ == "__main__":
+    import torch
+    from torch.autograd import Variable
+    from dynamic_loss_scaler import DynamicLossScaler
+
+    # N is batch size; D_in is input dimension;
+    # H is hidden dimension; D_out is output dimension.
+    N, D_in, H, D_out = 64, 1000, 100, 10
+
+    # Create random Tensors to hold inputs and outputs, and wrap them in Variables.
+    x = Variable(torch.randn(N, D_in), requires_grad=False)
+    y = Variable(torch.randn(N, D_out), requires_grad=False)
+
+    w1 = Variable(torch.randn(D_in, H), requires_grad=True)
+    w2 = Variable(torch.randn(H, D_out), requires_grad=True)
+    parameters = [w1, w2]
+
+    learning_rate = 1e-6
+    optimizer = torch.optim.SGD(parameters, lr=learning_rate)
+    loss_scaler = DynamicLossScaler()
+
+    for t in range(500):
+        y_pred = x.mm(w1).clamp(min=0).mm(w2)
+        loss = (y_pred - y).pow(2).sum() * loss_scaler.loss_scale
+        print('Iter {} loss scale: {}'.format(t, loss_scaler.loss_scale))
+        print('Iter {} scaled loss: {}'.format(t, loss.data[0]))
+        print('Iter {} unscaled loss: {}'.format(t, loss.data[0] / loss_scaler.loss_scale))
+
+        # Run backprop
+        optimizer.zero_grad()
+        loss.backward()
+        
+        # Check for overflow
+        has_overflow = DynamicLossScaler.has_overflow(parameters)
+        
+        # If no overflow, unscale grad and update as usual
+        if not has_overflow:
+            for param in parameters:
+                param.grad.data.mul_(1. / loss_scaler.loss_scale)
+            optimizer.step()
+        # Otherwise, don't do anything -- ie, skip iteration
+        else:
+            print('OVERFLOW!')
+
+        # Update loss scale for next iteration
+        loss_scaler.update_scale(has_overflow)
+
+"""
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fused_dense/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fused_dense/__init__.py
new file mode 100644
index 0000000000..83d12cab27
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fused_dense/__init__.py
@@ -0,0 +1 @@
+from .fused_dense import *
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fused_dense/fused_dense.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fused_dense/fused_dense.py
new file mode 100644
index 0000000000..d36078c948
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/fused_dense/fused_dense.py
@@ -0,0 +1,85 @@
+import torch
+from torch import nn
+import fused_dense_cuda
+from .. import amp
+#implements fused GEMM+bias in forward pass using mlp_cuda from apex
+class FusedDenseFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, weight, bias):
+        ctx.save_for_backward(input, weight)
+        output = fused_dense_cuda.linear_bias_forward(input, weight, bias)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, weight = ctx.saved_tensors
+        grad_input, grad_weight, grad_bias = fused_dense_cuda.linear_bias_backward(input, weight, grad_output)
+        return grad_input, grad_weight, grad_bias
+
+class DenseNoBiasFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, weight):
+        ctx.save_for_backward(input, weight)
+        output = torch.matmul(input, weight.t())
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, weight = ctx.saved_tensors
+        grad_input = grad_output.mm(weight)
+        grad_weight = grad_output.t().mm(input)
+        return grad_input, grad_weight
+
+
+class FusedDenseGeluDenseFunc(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, weight1, bias1, weight2, bias2):
+        ctx.save_for_backward(input, weight1, weight2)
+        output1, output2, gelu_in = fused_dense_cuda.linear_gelu_linear_forward(input, weight1, bias1, weight2, bias2)
+        ctx.save_for_backward(input, weight1, weight2, gelu_in, output1)
+        return output2
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, weight1, weight2, gelu_in, output1 = ctx.saved_tensors
+        grad_input, grad_weight1, grad_bias1, grad_weight2, grad_bias2 = fused_dense_cuda.linear_gelu_linear_backward(input, gelu_in, output1, weight1, weight2, grad_output)
+        return grad_input, grad_weight1, grad_bias1, grad_weight2, grad_bias2
+
+
+fused_dense_function = amp.half_function(FusedDenseFunc.apply)
+dense_no_bias_function = amp.half_function(DenseNoBiasFunc.apply)
+fused_dense_gelu_dense_function = amp.half_function(FusedDenseGeluDenseFunc.apply)
+
+class FusedDense(nn.Module):
+    def __init__(self, in_features, out_features, bias=True):
+        super(FusedDense, self).__init__()
+        self.in_features = in_features
+        self.out_features = out_features
+        self.weight = nn.Parameter(torch.Tensor(out_features, in_features))
+        if bias:
+            self.bias = nn.Parameter(torch.Tensor(out_features))
+        else:
+            #assert False, "no-bias option not added yet"
+            self.register_parameter('bias', None)
+
+    def forward(self, input):
+        if self.bias is not None:
+            return fused_dense_function(input, self.weight, self.bias)
+        else:
+            return dense_no_bias_function(input, self.weight)
+
+class FusedDenseGeluDense(nn.Module):
+    def __init__(self, in_features, intermediate_features, out_features, bias=True):
+        super(FusedDenseGeluDense, self).__init__()
+        assert bias == True, "DenseGeluDense module without bias is currently not supported"
+        self.in_features = in_features
+        self.intermediate_features = intermediate_features
+        self.out_features = out_features
+        self.weight1 = nn.Parameter(torch.Tensor(intermediate_features, in_features))
+        self.bias1 = nn.Parameter(torch.Tensor(intermediate_features))
+        self.weight2 = nn.Parameter(torch.Tensor(out_features, intermediate_features))
+        self.bias2 = nn.Parameter(torch.Tensor(out_features))
+
+    def forward(self, input):
+        return fused_dense_gelu_dense_function(input, self.weight1, self.bias1, self.weight2, self.bias2)
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/mlp/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/mlp/__init__.py
new file mode 100644
index 0000000000..f2f30f7d13
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/mlp/__init__.py
@@ -0,0 +1 @@
+from .mlp import *
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/mlp/mlp.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/mlp/mlp.py
new file mode 100644
index 0000000000..bae38f3f89
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/mlp/mlp.py
@@ -0,0 +1,79 @@
+from copy import copy
+import math
+import torch
+from torch import nn
+import mlp_cuda
+from .. import amp
+
+class MlpFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, bias, activation, *args):
+        output = mlp_cuda.forward(bias, activation, args)
+        ctx.save_for_backward(*args)
+        ctx.outputs = output
+        ctx.bias = bias
+        ctx.activation = activation
+        return output[0]
+
+    @staticmethod
+    def backward(ctx, grad_o):
+        grads = mlp_cuda.backward(ctx.bias, ctx.activation, grad_o, ctx.outputs, ctx.saved_tensors)
+        del ctx.outputs
+        return (None, None, *grads)
+
+mlp_function = amp.half_function(MlpFunction.apply)
+
+class MLP(torch.nn.Module):
+    """Launch MLP in C++
+
+    Args:
+        mlp_sizes (list of int): MLP sizes. Example: [1024,1024,1024] will create 2 MLP layers with shape 1024x1024
+        bias (bool): Default True:
+        relu (bool): Default True
+    """
+    def __init__(self, mlp_sizes, bias=True, activation='relu'):
+        super(MLP, self).__init__()
+        self.num_layers = len(mlp_sizes) - 1
+        self.mlp_sizes = copy(mlp_sizes)
+        self.bias = 1 if bias else 0
+
+        if activation is 'none':
+            self.activation = 0
+        elif activation is 'relu':
+            self.activation = 1
+        elif activation is 'sigmoid':
+            self.activation = 2
+        else:
+            raise TypeError("activation must be relu or none.")
+
+        self.weights = []
+        self.biases = []
+        for i in range(self.num_layers):
+            w = torch.nn.Parameter(torch.empty(mlp_sizes[i+1], mlp_sizes[i]))
+            self.weights.append(w)
+            name = 'weight_{}'.format(i)
+            setattr(self, name, w)
+            if self.bias:
+                b = torch.nn.Parameter(torch.empty(mlp_sizes[i+1]))
+                self.biases.append(b)
+                name = 'bias_{}'.format(i)
+                setattr(self, name, b)
+
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        for weight in self.weights:
+            dimsum = weight.size(0) + weight.size(1)
+            std = math.sqrt(2. / float(dimsum))
+            nn.init.normal_(weight, 0., std)
+        if self.bias:
+            for bias in self.biases:
+                std = math.sqrt(1. / float(bias.size(0)))
+                nn.init.normal_(bias, 0., std)
+
+    def forward(self, input):
+        return mlp_function(self.bias, self.activation, input, *self.weights, *self.biases)
+
+    def extra_repr(self):
+        s = F"MLP sizes: {self.mlp_sizes}, Bias={self.bias}, activation={self.activation}"
+        return s
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/multi_tensor_apply/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/multi_tensor_apply/__init__.py
new file mode 100644
index 0000000000..0a80e3c545
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/multi_tensor_apply/__init__.py
@@ -0,0 +1,4 @@
+from .multi_tensor_apply import MultiTensorApply
+
+multi_tensor_applier = MultiTensorApply(2048*32)
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/multi_tensor_apply/multi_tensor_apply.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/multi_tensor_apply/multi_tensor_apply.py
new file mode 100644
index 0000000000..346c6e50fb
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/multi_tensor_apply/multi_tensor_apply.py
@@ -0,0 +1,30 @@
+import torch
+
+class MultiTensorApply(object):
+    available = False
+    warned = False
+
+    def __init__(self, chunk_size):
+        try:
+            import amp_C
+            MultiTensorApply.available = True
+            self.chunk_size = chunk_size
+        except ImportError as err:
+            MultiTensorApply.available = False
+            MultiTensorApply.import_err = err
+
+    def check_avail(self):
+        if MultiTensorApply.available == False:
+            raise RuntimeError(
+                "Attempted to call MultiTensorApply method, but MultiTensorApply "
+                "is not available, possibly because Apex was installed without "
+                "--cpp_ext --cuda_ext.  Original import error message:",
+                MultiTensorApply.import_err)
+
+    def __call__(self, op, noop_flag_buffer, tensor_lists, *args):
+        self.check_avail()
+
+        return op(self.chunk_size,
+                  noop_flag_buffer,
+                  tensor_lists,
+                  *args)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/normalization/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/normalization/__init__.py
new file mode 100644
index 0000000000..07941f271a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/normalization/__init__.py
@@ -0,0 +1 @@
+from .fused_layer_norm import FusedLayerNorm, MixedFusedLayerNorm
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/normalization/fused_layer_norm.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/normalization/fused_layer_norm.py
new file mode 100644
index 0000000000..337af76a33
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/normalization/fused_layer_norm.py
@@ -0,0 +1,218 @@
+import importlib
+import numbers
+
+import torch
+from torch.nn.parameter import Parameter
+from torch.nn import init
+from torch.nn import functional as F
+
+from apex._autocast_utils import _cast_if_autocast_enabled
+
+global fused_layer_norm_cuda
+fused_layer_norm_cuda = None
+
+
+class FusedLayerNormAffineFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, weight, bias, normalized_shape, eps):
+        global fused_layer_norm_cuda
+        if fused_layer_norm_cuda is None:
+            fused_layer_norm_cuda = importlib.import_module("fused_layer_norm_cuda")
+        ctx.normalized_shape = normalized_shape
+        ctx.eps = eps
+        input_ = input.contiguous()
+        weight_ = weight.contiguous()
+        bias_ = bias.contiguous()
+        output, mean, invvar = fused_layer_norm_cuda.forward_affine(
+            input_, ctx.normalized_shape, weight_, bias_, ctx.eps
+        )
+        ctx.save_for_backward(input_, weight_, bias_, mean, invvar)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input_, weight_, bias_, mean, invvar = ctx.saved_tensors
+        grad_input = grad_weight = grad_bias = None
+        grad_input, grad_weight, grad_bias = fused_layer_norm_cuda.backward_affine(
+            grad_output.contiguous(), mean, invvar, input_, ctx.normalized_shape, weight_, bias_, ctx.eps
+        )
+        return grad_input, grad_weight, grad_bias, None, None
+
+
+class FusedLayerNormAffineMixedDtypesFunction(FusedLayerNormAffineFunction):
+
+    @staticmethod
+    def forward(ctx, input, weight, bias, normalized_shape, eps):
+        global fused_layer_norm_cuda
+        if fused_layer_norm_cuda is None:
+            fused_layer_norm_cuda = importlib.import_module("fused_layer_norm_cuda")
+        ctx.normalized_shape = normalized_shape
+        ctx.eps = eps
+        input_ = input.contiguous()
+        weight_ = weight.contiguous()
+        bias_ = bias.contiguous()
+        output, mean, invvar = fused_layer_norm_cuda.forward_affine_mixed_dtypes(
+            input_, ctx.normalized_shape, weight_, bias_, ctx.eps
+        )
+        ctx.save_for_backward(input_, weight_, bias_, mean, invvar)
+        return output
+
+
+class FusedLayerNormFunction(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, input, normalized_shape, eps):
+        global fused_layer_norm_cuda
+        if fused_layer_norm_cuda is None:
+            fused_layer_norm_cuda = importlib.import_module("fused_layer_norm_cuda")
+        ctx.normalized_shape = normalized_shape
+        ctx.eps = eps
+        input_ = input.contiguous()
+        output, mean, invvar = fused_layer_norm_cuda.forward(input_, ctx.normalized_shape, ctx.eps)
+        ctx.save_for_backward(input_, mean, invvar)
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input_, mean, invvar = ctx.saved_tensors
+        grad_input = None
+        grad_input = fused_layer_norm_cuda.backward(
+            grad_output.contiguous(), mean, invvar, input_, ctx.normalized_shape, ctx.eps
+        )
+        return grad_input, None, None
+
+
+def fused_layer_norm_affine(input, weight, bias, normalized_shape, eps=1e-6):
+    args = _cast_if_autocast_enabled(input, weight, bias, normalized_shape, eps)
+    with torch.cuda.amp.autocast(enabled=False):
+        return FusedLayerNormAffineFunction.apply(*args)
+
+
+def fused_layer_norm(input, normalized_shape, eps=1e-6):
+    args = _cast_if_autocast_enabled(input, normalized_shape, eps)
+    with torch.cuda.amp.autocast(enabled=False):
+        return FusedLayerNormFunction.apply(*args)
+
+
+def mixed_dtype_fused_layer_norm_affine(input, weight, bias, normalized_shape, eps=1e-6):
+    args = _cast_if_autocast_enabled(input, weight, bias, normalized_shape, eps)
+    with torch.cuda.amp.autocast(enabled=False):
+        return FusedLayerNormAffineMixedDtypesFunction.apply(*args)
+
+
+class FusedLayerNorm(torch.nn.Module):
+    r"""Applies Layer Normalization over a mini-batch of inputs as described in
+    the paper `Layer Normalization`_ .
+
+    Currently only runs on cuda() tensors.
+
+    .. math::
+        y = \frac{x - \mathrm{E}[x]}{ \sqrt{\mathrm{Var}[x] + \epsilon}} * \gamma + \beta
+
+    The mean and standard-deviation are calculated separately over the last
+    certain number dimensions which have to be of the shape specified by
+    :attr:`normalized_shape`.
+    :math:`\gamma` and :math:`\beta` are learnable affine transform parameters of
+    :attr:`normalized_shape` if :attr:`elementwise_affine` is ``True``.
+
+    .. note::
+        Unlike Batch Normalization and Instance Normalization, which applies
+        scalar scale and bias for each entire channel/plane with the
+        :attr:`affine` option, Layer Normalization applies per-element scale and
+        bias with :attr:`elementwise_affine`.
+
+    This layer uses statistics computed from input data in both training and
+    evaluation modes.
+
+    Args:
+        normalized_shape (int or list or torch.Size): input shape from an expected input
+            of size
+
+            .. math::
+                [* \times \text{normalized}\_\text{shape}[0] \times \text{normalized}\_\text{shape}[1]
+                    \times \ldots \times \text{normalized}\_\text{shape}[-1]]
+
+            If a single integer is used, it is treated as a singleton list, and this module will
+            normalize over the last dimension which is expected to be of that specific size.
+        eps: a value added to the denominator for numerical stability. Default: 1e-5
+        elementwise_affine: a boolean value that when set to ``True``, this module
+            has learnable per-element affine parameters initialized to ones (for weights)
+            and zeros (for biases). Default: ``True``.
+
+    Shape:
+        - Input: :math:`(N, *)`
+        - Output: :math:`(N, *)` (same shape as input)
+
+    Examples::
+
+        >>> input = torch.randn(20, 5, 10, 10)
+        >>> # With Learnable Parameters
+        >>> m = apex.normalization.FusedLayerNorm(input.size()[1:])
+        >>> # Without Learnable Parameters
+        >>> m = apex.normalization.FusedLayerNorm(input.size()[1:], elementwise_affine=False)
+        >>> # Normalize over last two dimensions
+        >>> m = apex.normalization.FusedLayerNorm([10, 10])
+        >>> # Normalize over last dimension of size 10
+        >>> m = apex.normalization.FusedLayerNorm(10)
+        >>> # Activating the module
+        >>> output = m(input)
+
+    .. _`Layer Normalization`: https://arxiv.org/abs/1607.06450
+    """
+
+    def __init__(self, normalized_shape, eps=1e-5, elementwise_affine=True):
+        super().__init__()
+
+        global fused_layer_norm_cuda
+        fused_layer_norm_cuda = importlib.import_module("fused_layer_norm_cuda")
+
+        if isinstance(normalized_shape, numbers.Integral):
+            normalized_shape = (normalized_shape,)
+        self.normalized_shape = torch.Size(normalized_shape)
+        self.eps = eps
+        self.elementwise_affine = elementwise_affine
+        if self.elementwise_affine:
+            self.weight = Parameter(torch.Tensor(*normalized_shape))
+            self.bias = Parameter(torch.Tensor(*normalized_shape))
+        else:
+            self.register_parameter("weight", None)
+            self.register_parameter("bias", None)
+        self.reset_parameters()
+
+    def reset_parameters(self):
+        if self.elementwise_affine:
+            init.ones_(self.weight)
+            init.zeros_(self.bias)
+
+    def forward(self, input):
+        if not input.is_cuda:
+            return F.layer_norm(input, self.normalized_shape, self.weight, self.bias, self.eps)
+        if self.elementwise_affine:
+            return fused_layer_norm_affine(input, self.weight, self.bias, self.normalized_shape, self.eps)
+        else:
+            return fused_layer_norm(input, self.normalized_shape, self.eps)
+
+    def extra_repr(self):
+        return "{normalized_shape}, eps={eps}, " "elementwise_affine={elementwise_affine}".format(**self.__dict__)
+
+
+# NOTE (mkozuki): Why "mixed"?
+# MixedFusedLayerNorm differs from FusedLayerNorm in that this layer norm uses parameter's dtype
+# as output tensor's dtype while FusedLayerNorm uses input tensor's dtype for output tensor's dtype.
+# See: `layer_norm_affine` and `layer_norm_affine_mixed_dtypes` in "csrc/layer_norm_cuda.cpp"
+class MixedFusedLayerNorm(FusedLayerNorm):
+
+    def __init__(self, normalized_shape, eps=1e-5, **kwargs):
+        if "elementwise_affine" in kwargs:
+            import warnings
+            warnings.warn("MixedFusedLayerNorm does not support `elementwise_affine` argument")
+            elementwise_affine = kwargs.pop("elementwise_affine")
+            if not elementwise_affine:
+                raise RuntimeError("MixedFusedLayerNorm does not support `elementwise_affine = False`")
+
+        super().__init__(normalized_shape=normalized_shape, eps=eps, elementwise_affine=True)
+
+    def forward(self, input: torch.Tensor):
+        # NOTE (mkozuki): CPU path is here mainly for unittest sake.
+        if not input.is_cuda:
+            return F.layer_norm(input, self.normalized_shape, self.weight, self.bias, self.eps)
+        return mixed_dtype_fused_layer_norm_affine(input, self.weight, self.bias, self.normalized_shape, self.eps)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/__init__.py
new file mode 100644
index 0000000000..52e8e274d8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/__init__.py
@@ -0,0 +1,5 @@
+from .fused_sgd import FusedSGD
+from .fused_adam import FusedAdam
+from .fused_novograd import FusedNovoGrad
+from .fused_lamb import FusedLAMB
+from .fused_adagrad import FusedAdagrad
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_adagrad.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_adagrad.py
new file mode 100644
index 0000000000..d72a68c5d2
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_adagrad.py
@@ -0,0 +1,122 @@
+import torch
+from apex.multi_tensor_apply import multi_tensor_applier
+
+
+class FusedAdagrad(torch.optim.Optimizer):
+    """Implements Adagrad algorithm.
+
+    Currently GPU-only.  Requires Apex to be installed via
+    ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``.
+
+    This version of fused Adagrad implements 2 fusions.
+      * Fusion of the Adagrad update's elementwise operations
+      * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
+
+    :class:`apex.optimizers.FusedAdagrad`'s usage is identical to any ordinary Pytorch optimizer::
+        opt = apex.optimizers.FusedAdagrad(model.parameters(), lr = ....)
+        ...
+        opt.step()
+
+    :class:`apex.optimizers.FusedAdagrad` may be used with or without Amp.  If you wish to use :class:`FusedAdagrad` with Amp,
+    you may choose any ``opt_level``::
+        opt = apex.optimizers.FusedAdagrad(model.parameters(), lr = ....)
+        model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
+        ...
+        opt.step()
+    In general, ``opt_level="O1"`` is recommended.
+
+    It has been proposed in `Adaptive Subgradient Methods for Online Learning
+    and Stochastic Optimization`_.
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-2)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-10)
+        adagrad_w_mode (boolean, optional): Apply L2 regularization or weight decay
+            True for decoupled weight decay (also known as AdamW) (default: False)
+
+    .. _Adaptive Subgradient Methods for Online Learning and Stochastic
+        Optimization: http://jmlr.org/papers/v12/duchi11a.html
+    """
+    def __init__(self, params, lr=1e-2, eps=1e-10,
+                 weight_decay=0., set_grad_none=True, adagrad_w_mode=False):
+
+        defaults = dict(lr=lr, eps=eps, weight_decay=weight_decay)
+        super(FusedAdagrad, self).__init__(params, defaults)
+        self.adagrad_w_mode = 1 if adagrad_w_mode else 0
+        self.set_grad_none = set_grad_none
+
+        if multi_tensor_applier.available:
+            import amp_C
+            # Skip buffer
+            self._dummy_overflow_buf = torch.cuda.IntTensor([0])
+            self.multi_tensor_adagrad = amp_C.multi_tensor_adagrad
+        else:
+            raise RuntimeError('apex.optimizers.FusedAdagrad requires cuda extensions')
+
+    def zero_grad(self):
+        if self.set_grad_none:
+            for group in self.param_groups:
+                for p in group['params']:
+                    p.grad = None
+        else:
+            super(FusedAdagrad, self).zero_grad()
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            # create lists for multi-tensor apply
+            g_16, p_16, h_16 = [], [], []
+            g_32, p_32, h_32 = [], [], []
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                if p.grad.data.is_sparse:
+                    raise RuntimeError('FusedAdagrad does not support sparse gradients')
+
+                state = self.state[p]
+                # State initialization
+                if len(state) == 0:
+                    # Exponential moving average of gradient values
+                    state['sum'] = torch.zeros_like(p.data)
+                if p.dtype == torch.float16:
+                    g_16.append(p.grad.data)
+                    p_16.append(p.data)
+                    h_16.append(state['sum'])
+                elif p.dtype == torch.float32:
+                    g_32.append(p.grad.data)
+                    p_32.append(p.data)
+                    h_32.append(state['sum'])
+                else:
+                    raise RuntimeError('FusedAdagrad only support fp16 and fp32.')
+
+            if(len(g_16) > 0):
+                multi_tensor_applier(self.multi_tensor_adagrad,
+                                     self._dummy_overflow_buf,
+                                     [g_16, p_16, h_16],
+                                     group['lr'],
+                                     group['eps'],
+                                     self.adagrad_w_mode,
+                                     group['weight_decay'])
+            if(len(g_32) > 0):
+                multi_tensor_applier(self.multi_tensor_adagrad,
+                                     self._dummy_overflow_buf,
+                                     [g_32, p_32, h_32],
+                                     group['lr'],
+                                     group['eps'],
+                                     self.adagrad_w_mode,
+                                     group['weight_decay'])
+
+        return loss
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_adam.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_adam.py
new file mode 100644
index 0000000000..742353bf22
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_adam.py
@@ -0,0 +1,173 @@
+import torch
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class FusedAdam(torch.optim.Optimizer):
+
+    """Implements Adam algorithm.
+
+    Currently GPU-only.  Requires Apex to be installed via
+    ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``.
+
+    This version of fused Adam implements 2 fusions.
+
+      * Fusion of the Adam update's elementwise operations
+      * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
+
+    :class:`apex.optimizers.FusedAdam` may be used as a drop-in replacement for ``torch.optim.AdamW``,
+    or ``torch.optim.Adam`` with ``adam_w_mode=False``::
+
+        opt = apex.optimizers.FusedAdam(model.parameters(), lr = ....)
+        ...
+        opt.step()
+
+    :class:`apex.optimizers.FusedAdam` may be used with or without Amp.  If you wish to use :class:`FusedAdam` with Amp,
+    you may choose any ``opt_level``::
+
+        opt = apex.optimizers.FusedAdam(model.parameters(), lr = ....)
+        model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
+        ...
+        opt.step()
+
+    In general, ``opt_level="O1"`` is recommended.
+
+
+    .. warning::
+        A previous version of :class:`FusedAdam` allowed a number of additional arguments to ``step``.  These additional arguments
+        are now deprecated and unnecessary.
+
+    Adam was been proposed in `Adam: A Method for Stochastic Optimization`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square. (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False) NOT SUPPORTED in FusedAdam!
+        adam_w_mode (boolean, optional): Apply L2 regularization or weight decay
+            True for decoupled weight decay(also known as AdamW) (default: True)
+        set_grad_none (bool, optional): whether set grad to None when zero_grad()
+            method is called. (default: True)
+
+    .. _Adam - A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params, lr=1e-3, bias_correction=True,
+                 betas=(0.9, 0.999), eps=1e-8, adam_w_mode=True,
+                 weight_decay=0., amsgrad=False, set_grad_none=True):
+
+        if amsgrad:
+            raise RuntimeError('FusedAdam does not support the AMSGrad variant.')
+        defaults = dict(lr=lr, bias_correction=bias_correction,
+                        betas=betas, eps=eps, weight_decay=weight_decay)
+        super(FusedAdam, self).__init__(params, defaults)
+        self.adam_w_mode = 1 if adam_w_mode else 0
+        self.set_grad_none = set_grad_none
+        if multi_tensor_applier.available:
+            import amp_C
+            # Skip buffer
+            self._dummy_overflow_buf = torch.cuda.IntTensor([0])
+            self.multi_tensor_adam = amp_C.multi_tensor_adam
+        else:
+            raise RuntimeError('apex.optimizers.FusedAdam requires cuda extensions')
+
+    def zero_grad(self):
+        if self.set_grad_none:
+            for group in self.param_groups:
+                for p in group['params']:
+                    p.grad = None
+        else:
+            super(FusedAdam, self).zero_grad()
+
+    def step(self, closure=None, grads=None, output_params=None, scale=None, grad_norms=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+
+        The remaining arguments are deprecated, and are only retained (for the moment) for error-checking purposes.
+        """
+        if any(p is not None for p in [grads, output_params, scale, grad_norms]):
+            raise RuntimeError('FusedAdam has been updated.  Simply initialize it identically to torch.optim.Adam, and call step() with no arguments.')
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            bias_correction = 1 if group['bias_correction'] else 0
+            beta1, beta2 = group['betas']
+
+            # assume same step across group now to simplify things
+            # per parameter step can be easily support by making it tensor, or pass list into kernel
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            # create lists for multi-tensor apply
+            g_16, p_16, m_16, v_16 = [], [], [], []
+            g_32, p_32, m_32, v_32 = [], [], [], []
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                if p.grad.data.is_sparse:
+                    raise RuntimeError('FusedAdam does not support sparse gradients, please consider SparseAdam instead')
+
+                state = self.state[p]
+                # State initialization
+                if len(state) == 0:
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+                if p.dtype == torch.float16:
+                    g_16.append(p.grad.data)
+                    p_16.append(p.data)
+                    m_16.append(state['exp_avg'])
+                    v_16.append(state['exp_avg_sq'])
+                elif p.dtype == torch.float32:
+                    g_32.append(p.grad.data)
+                    p_32.append(p.data)
+                    m_32.append(state['exp_avg'])
+                    v_32.append(state['exp_avg_sq'])
+                else:
+                    raise RuntimeError('FusedAdam only support fp16 and fp32.')
+
+            if(len(g_16) > 0):
+                multi_tensor_applier(self.multi_tensor_adam,
+                                     self._dummy_overflow_buf,
+                                     [g_16, p_16, m_16, v_16],
+                                     group['lr'],
+                                     beta1,
+                                     beta2,
+                                     group['eps'],
+                                     group['step'],
+                                     self.adam_w_mode,
+                                     bias_correction,
+                                     group['weight_decay'])
+            if(len(g_32) > 0):
+                multi_tensor_applier(self.multi_tensor_adam,
+                                     self._dummy_overflow_buf,
+                                     [g_32, p_32, m_32, v_32],
+                                     group['lr'],
+                                     beta1,
+                                     beta2,
+                                     group['eps'],
+                                     group['step'],
+                                     self.adam_w_mode,
+                                     bias_correction,
+                                     group['weight_decay'])
+
+
+        return loss
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_lamb.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_lamb.py
new file mode 100644
index 0000000000..854525dcfb
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_lamb.py
@@ -0,0 +1,215 @@
+import torch
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class FusedLAMB(torch.optim.Optimizer):
+
+    """Implements LAMB algorithm.
+
+    Currently GPU-only.  Requires Apex to be installed via
+    ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``.
+
+    This version of fused LAMB implements 2 fusions.
+
+      * Fusion of the LAMB update's elementwise operations
+      * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
+
+    :class:`apex.optimizers.FusedLAMB`'s usage is identical to any ordinary Pytorch optimizer::
+
+        opt = apex.optimizers.FusedLAMB(model.parameters(), lr = ....)
+        ...
+        opt.step()
+
+    :class:`apex.optimizers.FusedLAMB` may be used with or without Amp.  If you wish to use :class:`FusedLAMB` with Amp,
+    you may choose any ``opt_level``::
+
+        opt = apex.optimizers.FusedLAMB(model.parameters(), lr = ....)
+        model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
+        ...
+        opt.step()
+
+    In general, ``opt_level="O1"`` is recommended.
+
+    LAMB was proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its norm. (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            NOT SUPPORTED now! (default: False)
+        adam_w_mode (boolean, optional): Apply L2 regularization or weight decay
+            True for decoupled weight decay(also known as AdamW) (default: True)
+        grad_averaging (bool, optional): whether apply (1-beta2) to grad when
+            calculating running averages of gradient. (default: True)
+        set_grad_none (bool, optional): whether set grad to None when zero_grad()
+            method is called. (default: True)
+        max_grad_norm (float, optional): value used to clip global grad norm
+            (default: 1.0)
+        use_nvlamb (boolean, optional): Apply adaptive learning rate to 0.0
+            weight decay parameter (default: False)
+
+    .. _Large Batch Optimization for Deep Learning - Training BERT in 76 minutes:
+        https://arxiv.org/abs/1904.00962
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params, lr=1e-3, bias_correction=True,
+                 betas=(0.9, 0.999), eps=1e-6, weight_decay=0.01,
+                 amsgrad=False, adam_w_mode=True,
+                 grad_averaging=True, set_grad_none=True,
+                 max_grad_norm=1.0, use_nvlamb=False):
+        if amsgrad:
+            raise RuntimeError('FusedLAMB does not support the AMSGrad variant.')
+        defaults = dict(lr=lr, bias_correction=bias_correction,
+                        betas=betas, eps=eps, weight_decay=weight_decay,
+                        grad_averaging=grad_averaging,
+                        max_grad_norm=max_grad_norm)
+        super(FusedLAMB, self).__init__(params, defaults)
+        if multi_tensor_applier.available:
+            import amp_C
+            self.multi_tensor_l2norm=amp_C.multi_tensor_l2norm
+            # Skip buffer
+            self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device=self.param_groups[0]["params"][0].device)
+            self.multi_tensor_lamb = amp_C.multi_tensor_lamb
+        else:
+            raise RuntimeError('apex.optimizers.FusedLAMB requires cuda extensions')
+
+        self.adam_w_mode = 1 if adam_w_mode else 0
+        self.set_grad_none = set_grad_none
+        self.use_nvlamb = use_nvlamb
+
+    def zero_grad(self):
+        if self.set_grad_none:
+            for group in self.param_groups:
+                for p in group['params']:
+                    p.grad = None
+        else:
+            super(FusedLAMB, self).zero_grad()
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        # create separate grad lists for fp32 and fp16 params
+        g_all_32, g_all_16 = [], []
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                if p.dtype == torch.float32:
+                    g_all_32.append(p.grad.data)
+                elif p.dtype == torch.float16:
+                    g_all_16.append(p.grad.data)
+                else:
+                    raise RuntimeError('FusedLAMB only support fp16 and fp32.')
+
+        device = self.param_groups[0]["params"][0].device
+        g_norm_32, g_norm_16 = torch.zeros(1, device=device), torch.zeros(1, device=device)
+        # compute grad norm for two lists
+        if len(g_all_32) > 0:
+            g_norm_32 = multi_tensor_applier(self.multi_tensor_l2norm,
+                                             self._dummy_overflow_buf,
+                                             [g_all_32], False)[0]
+        if len(g_all_16) > 0:
+            g_norm_16 = multi_tensor_applier(self.multi_tensor_l2norm,
+                                             self._dummy_overflow_buf,
+                                             [g_all_16], False)[0]
+
+        # blend two grad norms to get global grad norm
+        global_grad_norm = multi_tensor_applier(self.multi_tensor_l2norm,
+                                                self._dummy_overflow_buf,
+                                                [[g_norm_32, g_norm_16]],
+                                                False)[0]
+        max_grad_norm = self.defaults['max_grad_norm']
+
+        for group in self.param_groups:
+            bias_correction = 1 if group['bias_correction'] else 0
+            beta1, beta2 = group['betas']
+            grad_averaging = 1 if group['grad_averaging'] else 0
+
+            # assume same step across group now to simplify things
+            # per parameter step can be easily support by making it tensor, or pass list into kernel
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            # create lists for multi-tensor apply
+            g_16, p_16, m_16, v_16 = [], [], [], []
+            g_32, p_32, m_32, v_32 = [], [], [], []
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                if p.grad.data.is_sparse:
+                    raise RuntimeError('FusedLAMB does not support sparse gradients, please consider SparseAdam instead')
+
+                state = self.state[p]
+                # State initialization
+                if len(state) == 0:
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+
+                if p.dtype == torch.float16:
+                    g_16.append(p.grad.data)
+                    p_16.append(p.data)
+                    m_16.append(state['exp_avg'])
+                    v_16.append(state['exp_avg_sq'])
+                elif p.dtype == torch.float32:
+                    g_32.append(p.grad.data)
+                    p_32.append(p.data)
+                    m_32.append(state['exp_avg'])
+                    v_32.append(state['exp_avg_sq'])
+                else:
+                    raise RuntimeError('FusedLAMB only support fp16 and fp32.')
+
+            if(len(g_16) > 0):
+                multi_tensor_applier(self.multi_tensor_lamb,
+                                     self._dummy_overflow_buf,
+                                     [g_16, p_16, m_16, v_16],
+                                     group['lr'],
+                                     beta1,
+                                     beta2,
+                                     group['eps'],
+                                     group['step'],
+                                     bias_correction,
+                                     group['weight_decay'],
+                                     grad_averaging,
+                                     self.adam_w_mode,
+                                     global_grad_norm,
+                                     max_grad_norm,
+                                     self.use_nvlamb)
+            if(len(g_32) > 0):
+                multi_tensor_applier(self.multi_tensor_lamb,
+                                     self._dummy_overflow_buf,
+                                     [g_32, p_32, m_32, v_32],
+                                     group['lr'],
+                                     beta1,
+                                     beta2,
+                                     group['eps'],
+                                     group['step'],
+                                     bias_correction,
+                                     group['weight_decay'],
+                                     grad_averaging,
+                                     self.adam_w_mode,
+                                     global_grad_norm,
+                                     max_grad_norm,
+                                     self.use_nvlamb)
+
+        return loss
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_novograd.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_novograd.py
new file mode 100644
index 0000000000..2820ae36c2
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_novograd.py
@@ -0,0 +1,214 @@
+import torch
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class FusedNovoGrad(torch.optim.Optimizer):
+
+    """Implements NovoGrad algorithm.
+
+    Currently GPU-only.  Requires Apex to be installed via
+    ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``.
+
+    This version of fused NovoGrad implements 2 fusions.
+
+      * Fusion of the NovoGrad update's elementwise operations
+      * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
+
+    :class:`apex.optimizers.FusedNovoGrad`'s usage is identical to any Pytorch optimizer::
+
+        opt = apex.optimizers.FusedNovoGrad(model.parameters(), lr = ....)
+        ...
+        opt.step()
+
+    :class:`apex.optimizers.FusedNovoGrad` may be used with or without Amp.  If you wish to use :class:`FusedNovoGrad` with Amp,
+    you may choose any ``opt_level``::
+
+        opt = apex.optimizers.FusedNovoGrad(model.parameters(), lr = ....)
+        model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
+        ...
+        opt.step()
+
+    In general, ``opt_level="O1"`` is recommended.
+
+    It has been proposed in `Jasper: An End-to-End Convolutional Neural Acoustic Model`_.
+    More info: https://nvidia.github.io/OpenSeq2Seq/html/optimizers.html#novograd
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups.
+        lr (float, optional): learning rate. (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its norm. (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability. (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            NOT SUPPORTED now! (default: False)
+        reg_inside_moment (bool, optional): whether do regularization (norm and L2)
+            in momentum calculation. True for include, False for not include and
+            only do it on update term. (default: False)
+        grad_averaging (bool, optional): whether apply (1-beta1) to grad when
+            calculating running averages of gradient. (default: True)
+        norm_type (int, optional): which norm to calculate for each layer.
+            2 for L2 norm, and 0 for infinite norm. These 2 are only supported
+            type now. (default: 2)
+        init_zero (bool, optional): whether init norm with 0 (start averaging on
+            1st step) or first step norm (start averaging on 2nd step). True for
+            init with 0. (default: False)
+        set_grad_none (bool, optional): whether set grad to None when zero_grad()
+            method is called. (default: True)
+
+    .. _Jasper - An End-to-End Convolutional Neural Acoustic Model:
+        https://arxiv.org/abs/1904.03288
+    .. _On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+
+    def __init__(self, params, lr=1e-3, bias_correction=True,
+                 betas=(0.9, 0.999), eps=1e-8, weight_decay=0.,
+                 amsgrad=False, reg_inside_moment=False,
+                 grad_averaging=True, norm_type=2, init_zero=False,
+                 set_grad_none=True):
+        if amsgrad:
+            raise RuntimeError('FusedNovoGrad does not support the AMSGrad variant.')
+        defaults = dict(lr=lr, bias_correction=bias_correction,
+                        betas=betas, eps=eps, weight_decay=weight_decay,
+                        grad_averaging=grad_averaging, norm_type=norm_type,
+                        init_zero=init_zero)
+        super(FusedNovoGrad, self).__init__(params, defaults)
+        if multi_tensor_applier.available:
+            import amp_C
+            # Skip buffer
+
+            # Creating the overflow buffer on the same device as the params tensors.
+            self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device=self.param_groups[0]["params"][0].device)
+            self.multi_tensor_novograd = amp_C.multi_tensor_novograd
+        else:
+            raise RuntimeError('apex.optimizers.FusedNovoGrad requires cuda extensions')
+
+        self.moment_mode = 0 if reg_inside_moment else 1
+        self.set_grad_none = set_grad_none
+
+    def zero_grad(self):
+        if self.set_grad_none:
+            for group in self.param_groups:
+                for p in group['params']:
+                    p.grad = None
+        else:
+            super(FusedNovoGrad, self).zero_grad()
+
+    def load_state_dict(self, state_dict):
+        super(FusedNovoGrad, self).load_state_dict(state_dict)
+        # in case exp_avg_sq is not on the same device as params, move it there
+        for group in self.param_groups:
+            if len(group['params']) > 0:
+                group['exp_avg_sq'][0] = group['exp_avg_sq'][0].to(group['params'][0].device)
+                group['exp_avg_sq'][1] = group['exp_avg_sq'][1].to(group['params'][0].device)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            bias_correction = 1 if group['bias_correction'] else 0
+            beta1, beta2 = group['betas']
+            grad_averaging = 1 if group['grad_averaging'] else 0
+
+            # assume same step across group now to simplify things
+            # per parameter step can be easily support by making it tensor, or pass list into kernel
+            if 'step' in group:
+                group['step'] += 1
+            else:
+                group['step'] = 1
+
+            # create lists for multi-tensor apply
+            g_16, p_16, m_16 = [], [], []
+            g_32, p_32, m_32 = [], [], []
+
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                if p.grad.data.is_sparse:
+                    raise RuntimeError('FusedNovoGrad does not support sparse gradients, please consider SparseAdam instead')
+
+                state = self.state[p]
+                # State initialization
+                if len(state) == 0:
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+
+                if p.dtype == torch.float16:
+                    g_16.append(p.grad.data)
+                    p_16.append(p.data)
+                    m_16.append(state['exp_avg'])
+                elif p.dtype == torch.float32:
+                    g_32.append(p.grad.data)
+                    p_32.append(p.data)
+                    m_32.append(state['exp_avg'])
+                else:
+                    raise RuntimeError('FusedNovoGrad only support fp16 and fp32.')
+
+            # we store per weight norm as one tensor for one group/precision combination
+            # different from optim.Adam, we store norm here(not ^2) so we can unify calculation for norm types
+            if 'exp_avg_sq' not in group:
+                group['exp_avg_sq'] = [None, None]
+                if group['init_zero']:
+                    # Creating the following parameters on the same device as the params tensors.
+                    group['exp_avg_sq'][0] = torch.cuda.FloatTensor(len(g_16), device=self.param_groups[0]["params"][0].device).contiguous().fill_(0)
+                    group['exp_avg_sq'][1] = torch.cuda.FloatTensor(len(g_32), device=self.param_groups[0]["params"][0].device).contiguous().fill_(0)
+                else: # init with first step norm, so first blend have no effect
+                    if group['norm_type'] == 0:
+                        v_16 = [torch.max(torch.abs(g.to(torch.float32))).item() for g in g_16]
+                        v_32 = [torch.max(torch.abs(g)).item() for g in g_32]
+                    elif group['norm_type'] == 2:
+                        v_16 = [torch.sum(torch.pow(g.to(torch.float32), 2)).sqrt().item() for g in g_16]
+                        v_32 = [torch.sum(torch.pow(g, 2)).sqrt().item() for g in g_32]
+                    else:
+                        raise RuntimeError('FusedNovoGrad only support l2/inf norm now.')
+                    # Creating the following parameters on the same device as the params tensors.
+                    group['exp_avg_sq'][0] = torch.cuda.FloatTensor(v_16, device=self.param_groups[0]["params"][0].device)
+                    group['exp_avg_sq'][1] = torch.cuda.FloatTensor(v_32, device=self.param_groups[0]["params"][0].device)
+            else:
+                assert(len(g_16) == group['exp_avg_sq'][0].numel())
+                assert(len(g_32) == group['exp_avg_sq'][1].numel())
+
+            if(len(g_16) > 0):
+                multi_tensor_applier(self.multi_tensor_novograd,
+                                     self._dummy_overflow_buf,
+                                     [g_16, p_16, m_16],
+                                     group['exp_avg_sq'][0],
+                                     group['lr'],
+                                     beta1,
+                                     beta2,
+                                     group['eps'],
+                                     group['step'],
+                                     bias_correction,
+                                     group['weight_decay'],
+                                     grad_averaging,
+                                     self.moment_mode,
+                                     group['norm_type'])
+            if(len(g_32) > 0):
+                multi_tensor_applier(self.multi_tensor_novograd,
+                                     self._dummy_overflow_buf,
+                                     [g_32, p_32, m_32],
+                                     group['exp_avg_sq'][1],
+                                     group['lr'],
+                                     beta1,
+                                     beta2,
+                                     group['eps'],
+                                     group['step'],
+                                     bias_correction,
+                                     group['weight_decay'],
+                                     grad_averaging,
+                                     self.moment_mode,
+                                     group['norm_type'])
+
+
+        return loss
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_sgd.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_sgd.py
new file mode 100644
index 0000000000..e7bdcb2b96
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/optimizers/fused_sgd.py
@@ -0,0 +1,227 @@
+import torch
+from torch.optim.optimizer import Optimizer, required
+
+from apex.multi_tensor_apply import multi_tensor_applier
+
+class FusedSGD(Optimizer):
+    r"""Implements stochastic gradient descent (optionally with momentum).
+
+    Currently GPU-only.  Requires Apex to be installed via
+    ``pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" ./``.
+
+    This version of fused SGD implements 2 fusions.
+
+      * Fusion of the SGD update's elementwise operations
+      * A multi-tensor apply launch that batches the elementwise updates applied to all the model's parameters into one or a few kernel launches.
+
+    :class:`apex.optimizers.FusedSGD` may be used as a drop-in replacement for ``torch.optim.SGD``::
+
+        opt = apex.optimizers.FusedSGD(model.parameters(), lr = ....)
+        ...
+        opt.step()
+
+    :class:`apex.optimizers.FusedSGD` may be used with or without Amp.  If you wish to use :class:`FusedSGD` with Amp,
+    you may choose any ``opt_level``::
+
+        opt = apex.optimizers.FusedSGD(model.parameters(), lr = ....)
+        model, opt = amp.initialize(model, opt, opt_level="O0" or "O1 or "O2")
+        ...
+        opt.step()
+
+    In general, ``opt_level="O1"`` is recommended.
+
+    Nesterov momentum is based on the formula from
+    `On the importance of initialization and momentum in deep learning`__.
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float): learning rate
+        momentum (float, optional): momentum factor (default: 0)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        dampening (float, optional): dampening for momentum (default: 0)
+        nesterov (bool, optional): enables Nesterov momentum (default: False)
+
+    Example:
+        >>> optimizer = torch.optim.SGD(model.parameters(), lr=0.1, momentum=0.9)
+        >>> optimizer.zero_grad()
+        >>> loss_fn(model(input), target).backward()
+        >>> optimizer.step()
+
+    __ http://www.cs.toronto.edu/%7Ehinton/absps/momentum.pdf
+
+    .. note::
+        The implementation of SGD with Momentum/Nesterov subtly differs from
+        Sutskever et. al. and implementations in some other frameworks.
+
+        Considering the specific case of Momentum, the update can be written as
+
+        .. math::
+                  v = \rho * v + g \\
+                  p = p - lr * v
+
+        where p, g, v and :math:`\rho` denote the parameters, gradient,
+        velocity, and momentum respectively.
+
+        This is in contrast to Sutskever et. al. and
+        other frameworks which employ an update of the form
+
+        .. math::
+             v = \rho * v + lr * g \\
+             p = p - v
+
+        The Nesterov version is analogously modified.
+    """
+
+    def __init__(self, params, lr=required, momentum=0, dampening=0,
+                 weight_decay=0, nesterov=False,
+                 wd_after_momentum=False,
+                 materialize_master_grads=True,
+                 set_grad_none=False):
+        if lr is not required and lr < 0.0:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if momentum < 0.0:
+            raise ValueError("Invalid momentum value: {}".format(momentum))
+        if weight_decay < 0.0:
+            raise ValueError("Invalid weight_decay value: {}".format(weight_decay))
+
+        defaults = dict(lr=lr, momentum=momentum, dampening=dampening,
+                        weight_decay=weight_decay, nesterov=nesterov)
+        if nesterov and (momentum <= 0 or dampening != 0):
+            raise ValueError("Nesterov momentum requires a momentum and zero dampening")
+        super(FusedSGD, self).__init__(params, defaults)
+
+        self.wd_after_momentum = wd_after_momentum
+        self.materialize_master_grads = materialize_master_grads
+        self.most_recent_scale = 1.0
+        self.scale_set_by_backward = False
+        self.set_grad_none = set_grad_none
+
+        if multi_tensor_applier.available:
+            import amp_C
+            # Skip buffer
+            self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device=self.param_groups[0]["params"][0].device)
+            self.multi_tensor_sgd = amp_C.multi_tensor_sgd
+        else:
+            raise RuntimeError('apex.optimizers.FusedSGD requires cuda extensions')
+
+    def __setstate__(self, state):
+        super(FusedSGD, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('nesterov', False)
+
+    def zero_grad(self):
+        if self.set_grad_none:
+            for group in self.param_groups:
+                for p in group['params']:
+                    p.grad = None
+        else:
+            super(FusedSGD, self).zero_grad()
+
+    def get_momentums(self, params):
+        momentums = []
+        first_run = True
+        for p in params:
+            param_state = self.state[p]
+            # torch.optim.SGD initializes momentum in the main loop, we have
+            # to do it here, and track whether or not we've done so, so that
+            # momentum application can be skipped in the main kernel.
+            if 'momentum_buffer' not in param_state:
+                first_run = True
+                buf = param_state['momentum_buffer'] = torch.zeros_like(p.data)
+                momentums.append(buf)
+            else:
+                first_run = False
+                momentums.append(param_state['momentum_buffer'])
+        return momentums, first_run
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        explicit_master_params = (hasattr(self, "_amp_stash") and
+                                  hasattr(self._amp_stash, "fp32_from_fp16_groups"))
+
+        for gid, group in enumerate(self.param_groups):
+            weight_decay = group['weight_decay']
+            momentum = group['momentum']
+            dampening = group['dampening']
+            nesterov = group['nesterov']
+
+
+            # For each group, there are 3 possible combinations we need to consider:
+            # grad_type, param_to_update_type, momentum_type, requires_fp16_model_copy
+            # 1. fp16, fp16, fp16, No
+            # 2. fp32, fp32, fp32, No
+            # 3. fp16, fp32, fp32, Yes
+
+            first_runs = [True, True]
+
+            # I think a bit of code divergence in exchange for naming clarity is worthwhile
+            if explicit_master_params:
+                stash = self._amp_stash
+
+                fp32_params = [p for p in stash.fp32_from_fp32_groups[gid] if p.grad is not None]
+                fp32_grads = [p.grad for p in stash.fp32_from_fp32_groups[gid] if p.grad is not None]
+                fp32_momentums, first_runs[1] = self.get_momentums(fp32_params)
+
+                if self.materialize_master_grads:
+                    fp16_model_params = [p for i, p in enumerate(
+                        stash.fp16_groups[gid]) if stash.fp32_from_fp16_groups[gid][i].grad is not None]
+                    fp32_from_fp16_grads = [p.grad for p in stash.fp32_from_fp16_groups[gid] if p.grad is not None]
+                    fp32_from_fp16_params = [p for p in stash.fp32_from_fp16_groups[gid] if p.grad is not None]
+                    fp32_from_fp16_momentums, first_runs[0] = self.get_momentums(fp32_from_fp16_params)
+
+                    fp16_set = [fp32_from_fp16_grads, fp32_from_fp16_params,
+                                fp32_from_fp16_momentums, fp16_model_params]
+                else:
+                    fp16_model_params = [p for p in stash.fp16_groups[gid] if p.grad is not None]
+                    fp16_model_grads = [p.grad for p in stash.fp16_groups[gid] if p.grad is not None]
+                    fp32_from_fp16_params = [p for i, p in enumerate(
+                        stash.fp32_from_fp16_groups[gid]) if stash.fp16_groups[gid][i].grad is not None]
+                    fp32_from_fp16_momentums, first_runs[0] = self.get_momentums(fp32_from_fp16_params)
+
+                    fp16_set = [fp16_model_grads, fp32_from_fp16_params,
+                                fp32_from_fp16_momentums, fp16_model_params]
+
+                launch_sets= [fp16_set, [fp32_grads, fp32_params, fp32_momentums]]
+            else:
+                fp16_params = [p for p in group['params'] if (p.dtype == torch.float16 and p.grad is not None)]
+                fp16_grads = [p.grad for p in group['params'] if (p.dtype == torch.float16 and p.grad is not None)]
+                fp16_momentums, first_runs[0] = self.get_momentums(fp16_params)
+
+                fp32_params = [p for p in group['params'] if (p.dtype == torch.float32 and p.grad is not None)]
+                fp32_grads = [p.grad for p in group['params'] if (p.dtype == torch.float32 and p.grad is not None)]
+                fp32_momentums, first_runs[1] = self.get_momentums(fp32_params)
+
+                launch_sets = [[fp16_grads, fp16_params, fp16_momentums],
+                               [fp32_grads, fp32_params, fp32_momentums]]
+
+            for s, (launch_set, first_run) in enumerate(zip(launch_sets, first_runs)):
+                assert len(launch_set[0]) == len(launch_set[1])
+                assert len(launch_set[0]) == len(launch_set[2])
+                if len(launch_set[0]) > 0:
+                    multi_tensor_applier(
+                        self.multi_tensor_sgd,
+                        self._dummy_overflow_buf,
+                        launch_set,
+                        weight_decay,
+                        momentum,
+                        dampening,
+                        group['lr'],
+                        nesterov,
+                        first_run,
+                        self.wd_after_momentum,
+                        1.0/self.most_recent_scale)
+
+        self.most_recent_scale = 1.0
+        self.scale_set_by_backward = False
+
+        return loss
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/LARC.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/LARC.py
new file mode 100644
index 0000000000..4a93fcd65b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/LARC.py
@@ -0,0 +1,107 @@
+import torch
+from torch import nn
+from torch.nn.parameter import Parameter
+
+class LARC(object):
+    """
+    :class:`LARC` is a pytorch implementation of both the scaling and clipping variants of LARC,
+    in which the ratio between gradient and parameter magnitudes is used to calculate an adaptive 
+    local learning rate for each individual parameter. The algorithm is designed to improve
+    convergence of large batch training.
+     
+    See https://arxiv.org/abs/1708.03888 for calculation of the local learning rate.
+
+    In practice it modifies the gradients of parameters as a proxy for modifying the learning rate
+    of the parameters. This design allows it to be used as a wrapper around any torch.optim Optimizer.
+
+    ```
+    model = ...
+    optim = torch.optim.Adam(model.parameters(), lr=...)
+    optim = LARC(optim)
+    ```
+
+    It can even be used in conjunction with apex.fp16_utils.FP16_optimizer.
+
+    ```
+    model = ...
+    optim = torch.optim.Adam(model.parameters(), lr=...)
+    optim = LARC(optim)
+    optim = apex.fp16_utils.FP16_Optimizer(optim)
+    ```
+
+    Args:
+        optimizer: Pytorch optimizer to wrap and modify learning rate for.
+        trust_coefficient: Trust coefficient for calculating the lr. See https://arxiv.org/abs/1708.03888
+        clip: Decides between clipping or scaling mode of LARC. If `clip=True` the learning rate is set to `min(optimizer_lr, local_lr)` for each parameter. If `clip=False` the learning rate is set to `local_lr*optimizer_lr`.
+        eps: epsilon kludge to help with numerical stability while calculating adaptive_lr
+    """
+
+    def __init__(self, optimizer, trust_coefficient=0.02, clip=True, eps=1e-8):
+        self.optim = optimizer
+        self.trust_coefficient = trust_coefficient
+        self.eps = eps
+        self.clip = clip
+
+    def __getstate__(self):
+        return self.optim.__getstate__()
+
+    def __setstate__(self, state):
+        self.optim.__setstate__(state)
+
+    @property
+    def state(self):
+        return self.optim.state
+
+    def __repr__(self):
+        return self.optim.__repr__()
+
+    @property
+    def param_groups(self):
+        return self.optim.param_groups
+
+    @param_groups.setter
+    def param_groups(self, value):
+        self.optim.param_groups = value
+    
+    def state_dict(self):
+        return self.optim.state_dict()
+
+    def load_state_dict(self, state_dict):
+        self.optim.load_state_dict(state_dict)
+
+    def zero_grad(self):
+        self.optim.zero_grad()
+
+    def add_param_group(self, param_group):
+        self.optim.add_param_group( param_group)
+
+    def step(self):
+        with torch.no_grad():
+            weight_decays = []
+            for group in self.optim.param_groups:
+                # absorb weight decay control from optimizer
+                weight_decay = group['weight_decay'] if 'weight_decay' in group else 0
+                weight_decays.append(weight_decay)
+                group['weight_decay'] = 0
+                for p in group['params']:
+                    if p.grad is None:
+                        continue
+                    param_norm = torch.norm(p.data)
+                    grad_norm = torch.norm(p.grad.data)
+
+                    if param_norm != 0 and grad_norm != 0:
+                        # calculate adaptive lr + weight decay
+                        adaptive_lr = self.trust_coefficient * (param_norm) / (grad_norm + param_norm * weight_decay + self.eps)
+
+                        # clip learning rate for LARC
+                        if self.clip:
+                            # calculation of adaptive_lr so that when multiplied by lr it equals `min(adaptive_lr, lr)`
+                            adaptive_lr = min(adaptive_lr/group['lr'], 1)
+
+                        p.grad.data += weight_decay * p.data
+                        p.grad.data *= adaptive_lr
+
+        self.optim.step()
+        # return weight decay control to optimizer
+        for i, group in enumerate(self.optim.param_groups):
+            group['weight_decay'] = weight_decays[i]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/README.md
new file mode 100644
index 0000000000..e7910d82f3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/README.md
@@ -0,0 +1,66 @@
+## Distributed Data Parallel
+
+distributed.py contains the source code for `apex.parallel.DistributedDataParallel`, a module wrapper that enables multi-process multi-GPU data parallel training optimized for NVIDIA's NCCL communication library.
+
+`apex.parallel.DistributedDataParallel` achieves high performance by overlapping communication with
+computation in the backward pass and bucketing smaller transfers to reduce the total number of
+transfers required.
+
+multiproc.py contains the source code for `apex.parallel.multiproc`, a launch utility that places one process on each of the node's available GPUs.
+
+#### [API Documentation](https://nvidia.github.io/apex/parallel.html)
+
+#### [Example/Walkthrough](https://github.com/NVIDIA/apex/tree/master/examples/distributed)
+
+#### [Imagenet example with Mixed Precision](https://github.com/NVIDIA/apex/tree/master/examples/imagenet)
+
+#### [Simple example with FP16_Optimizer](https://github.com/NVIDIA/apex/tree/master/examples/FP16_Optimizer_simple/distributed_apex)
+
+### Synchronized Batch Normalization
+
+`apex.parallel.SyncBatchNorm` has similar APIs as with `torch.nn.BatchNorm*N*d`.
+It reduces stats on the first (channel) dimension of the Tensor and accepts
+arbitrary spatial dimensions.
+
+#### Installation
+
+Apex provides two sync BN implementation:
+
+1. There is the Python-only implementation, which is the default implementation
+when install with `python setup.py install`.
+It uses PyTorch primitive operations and distributed communication package from
+`torch.distributed`.
+
+   - _Python-only implementation requires input tensor to be of same data type as
+layer_
+
+2. We also provide implementation with kernels through CUDA/C++ extension with
+improved performance. We are experimenting with Welford and Kahan for reduction
+hoping to get better accuracy.
+   To use the kernel implementation, user need to install Apex with CUDA extension
+enabled `python setup.py install --cuda_ext`.
+
+   - _Custom kernel implementation supports fp16 input with fp32 layer as cudnn.
+This is required to run imagenet example in fp16._
+
+   - _Currently kernel implementation only supports GPU._
+
+#### HowTo
+
+1. User could use `apex.parallel.SyncBatchNorm` by building their module with
+the layer explicitly.
+
+```
+import apex
+input_t = torch.randn(3, 5, 20).cuda()
+sbn = apex.parallel.SyncBatchNorm(5).cuda()
+output_t = sbn(input)
+```
+
+2. User could also take a constructed `torch.nn.Model` and replace all its `torch.nn.BatchNorm*N*d` modules with `apex.parallel.SyncBatchNorm` through utility function `apex.parallel.convert_syncbn_model`.
+
+```
+# model is an instance of torch.nn.Module
+import apex
+sync_bn_model = apex.parallel.convert_syncbn_model(model)
+```
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/__init__.py
new file mode 100644
index 0000000000..3cd7ae56e8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/__init__.py
@@ -0,0 +1,95 @@
+import torch
+
+if hasattr(torch.distributed, 'ReduceOp'):
+    ReduceOp = torch.distributed.ReduceOp
+elif hasattr(torch.distributed, 'reduce_op'):
+    ReduceOp = torch.distributed.reduce_op
+else:
+    ReduceOp = torch.distributed.deprecated.reduce_op
+
+from .distributed import DistributedDataParallel, Reducer
+# This is tricky because I'd like SyncBatchNorm to be exposed the same way
+# for both the cuda-enabled and python-fallback versions, and I don't want
+# to suppress the error information.
+try:
+    import syncbn
+    from .optimized_sync_batchnorm import SyncBatchNorm
+except ImportError as err:
+    from .sync_batchnorm import SyncBatchNorm
+    SyncBatchNorm.syncbn_import_error = err
+
+def convert_syncbn_model(module, process_group=None, channel_last=False):
+    '''
+    Recursively traverse module and its children to replace all instances of
+    ``torch.nn.modules.batchnorm._BatchNorm`` with :class:`apex.parallel.SyncBatchNorm`.
+
+    All ``torch.nn.BatchNorm*N*d`` wrap around
+    ``torch.nn.modules.batchnorm._BatchNorm``, so this function lets you easily switch
+    to use sync BN.
+
+    Args:
+        module (torch.nn.Module): input module
+
+    Example::
+
+        >>> # model is an instance of torch.nn.Module
+        >>> import apex
+        >>> sync_bn_model = apex.parallel.convert_syncbn_model(model)
+    '''
+    mod = module
+    if isinstance(module, torch.nn.modules.instancenorm._InstanceNorm):
+        return module
+    if isinstance(module, torch.nn.modules.batchnorm._BatchNorm):
+        mod = SyncBatchNorm(module.num_features, module.eps, module.momentum, module.affine, module.track_running_stats, process_group, channel_last=channel_last)
+        mod.running_mean = module.running_mean
+        mod.running_var = module.running_var
+        mod.num_batches_tracked = module.num_batches_tracked
+        if module.affine:
+            mod.weight.data = module.weight.data.clone().detach()
+            mod.bias.data = module.bias.data.clone().detach()
+    for name, child in module.named_children():
+        mod.add_module(name, convert_syncbn_model(child,
+                                                  process_group=process_group,
+                                                  channel_last=channel_last))
+    # TODO(jie) should I delete model explicitly?
+    del module
+    return mod
+
+def create_syncbn_process_group(group_size):
+    '''
+    Creates process groups to be used for syncbn of a give ``group_size`` and returns
+    process group that current GPU participates in.
+
+    ``group_size`` must divide the total number of GPUs (world_size).
+
+    ``group_size`` of 0 would be considered as =world_size. In this case ``None`` will be returned.
+
+    ``group_size`` of 1 would be equivalent to using non-sync bn, but will still carry the overhead.
+
+    Args:
+        group_size (int): number of GPU's to collaborate for sync bn
+
+    Example::
+
+        >>> # model is an instance of torch.nn.Module
+        >>> import apex
+        >>> group = apex.parallel.create_syncbn_process_group(group_size)
+    '''
+
+    if group_size==0:
+        return None
+
+    world_size = torch.distributed.get_world_size()
+    assert(world_size >= group_size)
+    assert(world_size % group_size == 0)
+
+    group=None
+    for group_num in (range(world_size//group_size)):
+        group_ids = range(group_num*group_size, (group_num+1)*group_size)
+        cur_group = torch.distributed.new_group(ranks=group_ids)
+        if (torch.distributed.get_rank()//group_size == group_num):
+            group = cur_group
+            #can not drop out and return here, every process must go through creation of all subgroups
+
+    assert(group is not None)
+    return group
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/distributed.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/distributed.py
new file mode 100644
index 0000000000..5267c834a1
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/distributed.py
@@ -0,0 +1,639 @@
+import torch
+import torch.distributed as dist
+from torch.nn.modules import Module
+from torch.autograd import Variable
+from collections import OrderedDict
+from itertools import chain
+import copy
+import importlib
+from ..multi_tensor_apply import multi_tensor_applier
+
+imported_flatten_impl = False
+
+def import_flatten_impl():
+    global flatten_impl, unflatten_impl, imported_flatten_impl
+    try:
+        import apex_C
+        flatten_impl = apex_C.flatten
+        unflatten_impl = apex_C.unflatten
+    except ImportError:
+        print("Warning:  apex was installed without --cpp_ext.  Falling back to Python flatten and unflatten.")
+        flatten_impl = torch._utils._flatten_dense_tensors
+        unflatten_impl = torch._utils._unflatten_dense_tensors
+    imported_flatten_impl = True
+
+def flatten(bucket):
+    if not imported_flatten_impl:
+        import_flatten_impl()
+    return flatten_impl(bucket)
+
+def unflatten(coalesced, bucket):
+    if not imported_flatten_impl:
+        import_flatten_impl()
+    return unflatten_impl(coalesced, bucket)
+
+# apply_dist_call requires that tensors in 'bucket' are all the same type.
+def apply_flat_dist_call(bucket, call, extra_args=None):
+
+    coalesced = flatten(bucket)
+
+    if extra_args is not None:
+        call(coalesced, *extra_args)
+    else:
+        call(coalesced)
+
+    if call is dist.all_reduce:
+        coalesced /= dist.get_world_size()
+
+    for buf, synced in zip(bucket, unflatten(coalesced, bucket)):
+        buf.copy_(synced)
+
+def split_half_float_double(tensors):
+    dtypes = ["torch.cuda.HalfTensor",  "torch.cuda.FloatTensor", "torch.cuda.DoubleTensor"]
+    buckets = []
+    for i, dtype in enumerate(dtypes):
+        bucket = [t for t in tensors if t.type() == dtype]
+        if bucket:
+            buckets.append(bucket)
+    return buckets
+
+def split_by_type(tensors):
+    buckets = OrderedDict()
+    for tensor in tensors:
+        tp = tensor.type()
+        if tp not in buckets:
+            buckets[tp] = []
+        buckets[tp].append(tensor)
+    return buckets
+
+# flat_dist_call organizes 'tensors' by type.
+def flat_dist_call(tensors, call, extra_args=None):
+    buckets = split_by_type(tensors)
+
+    for tp in buckets:
+        bucket = buckets[tp]
+        apply_flat_dist_call(bucket, call, extra_args)
+
+
+def extract_tensors(maybe_tensor, tensor_list):
+    if torch.is_tensor(maybe_tensor):
+        tensor_list.append(maybe_tensor)
+    else:
+        try:
+            for item in maybe_tensor:
+                extract_tensors(item, tensor_list)
+        except TypeError:
+            return
+
+
+class Reducer(object):
+    """
+    :class:`apex.parallel.Reducer` is a simple class that helps allreduce a module's parameters
+    across processes.  :class:`Reducer` is intended to give the user additional control:
+    Unlike :class:`DistributedDataParallel`, :class:`Reducer` will not automatically allreduce
+    parameters during ``backward()``.
+    Instead, :class:`Reducer` waits for the user to call ``<reducer_instance>.reduce()`` manually.
+    This enables, for example, delaying the allreduce to be carried out every
+    several iterations instead of every single iteration.
+
+    Like :class:`DistributedDataParallel`, :class:`Reducer` averages any tensors it allreduces
+    over the number of participating processes.
+
+    :class:`Reducer` is designed to work with the upstream launch utility script
+    ``torch.distributed.launch`` with ``--nproc_per_node <= number of gpus per node``.
+    When used with this launcher, :class:`Reducer` assumes 1:1 mapping of processes to GPUs.
+    It also assumes that your script calls ``torch.cuda.set_device(args.rank)`` before creating the model.
+
+    Args:
+        module_or_grads_list: Either a network definition (module) being run in multi-gpu/distributed mode, or an iterable of gradients to be reduced.  If a module is passed in, the Reducer constructor will sync the parameters across processes (broadcasting from rank 0) to make sure they're all initialized with the same values.  If a list of gradients (that came from some module) is passed in, the user is responsible for manually syncing that module's parameters at the beginning of training.
+    """
+
+    def __init__(self, module_or_grads_list):
+        if isinstance(module_or_grads_list, Module):
+            self.module = module_or_grads_list
+            flat_dist_call([param.data for param in self.module.parameters()], dist.broadcast, (0,) )
+
+        else:
+            self.module = None
+            self.grads = []
+            extract_tensors(module_or_grads_list, self.grads)
+
+    def reduce(self):
+        if self.module:
+            grads = [param.grad.data for param in self.module.parameters() if param.grad is not None]
+            flat_dist_call(grads, dist.all_reduce)
+        else:
+            flat_dist_call(self.grads, dist.all_reduce)
+
+
+class DistributedDataParallel(Module):
+    """
+    :class:`apex.parallel.DistributedDataParallel` is a module wrapper that enables
+    easy multiprocess distributed data parallel training, similar to ``torch.nn.parallel.DistributedDataParallel``.  Parameters are broadcast across participating processes on initialization, and gradients are
+    allreduced and averaged over processes during ``backward()``.
+
+    :class:`DistributedDataParallel` is optimized for use with NCCL.  It achieves high performance by
+    overlapping communication with computation during ``backward()`` and bucketing smaller gradient
+    transfers to reduce the total number of transfers required.
+
+    :class:`DistributedDataParallel` is designed to work with the upstream launch utility script
+    ``torch.distributed.launch`` with ``--nproc_per_node <= number of gpus per node``.
+    When used with this launcher, :class:`DistributedDataParallel` assumes 1:1 mapping of processes to GPUs.
+    It also assumes that your script calls ``torch.cuda.set_device(args.rank)`` before creating the model.
+
+    https://github.com/NVIDIA/apex/tree/master/examples/simple/distributed shows detailed usage.
+    https://github.com/NVIDIA/apex/tree/master/examples/imagenet shows another example
+    that combines :class:`DistributedDataParallel` with mixed precision training.
+
+    Args:
+        module: Network definition to be run in multi-gpu/distributed mode.
+        message_size (int, default=1e7): Minimum number of elements in a communication bucket.
+        delay_allreduce (bool, default=False):  Delay all communication to the end of the backward pass.  This disables overlapping communication with computation.
+        allreduce_trigger_params (list, optional, default=None):  If supplied, should contain a list of parameters drawn from the model.  Allreduces will be kicked off whenever one of these parameters receives its gradient (as opposed to when a bucket of size message_size is full).  At the end of backward(), a cleanup allreduce to catch any remaining gradients will also be performed automatically.  If allreduce_trigger_params is supplied, the message_size argument will be ignored.
+        allreduce_always_fp32 (bool, default=False):  Convert any FP16 gradients to FP32 before allreducing.  This can improve stability for widely scaled-out runs.
+        gradient_average (bool, default=True):  Option to toggle whether or not DDP averages the allreduced gradients over processes.  For proper scaling, the default value of True is recommended.
+        gradient_predivide_factor (float, default=1.0):  Allows perfoming the average of gradients over processes partially before and partially after the allreduce.  Before allreduce:  ``grads.mul_(1.0/gradient_predivide_factor)``.  After allreduce:  ``grads.mul_(gradient_predivide_factor/world size)``.  This can reduce the stress on the dynamic range of FP16 allreduces for widely scaled-out runs.
+
+    .. warning::
+        If ``gradient_average=False``, the pre-allreduce division (``grads.mul_(1.0/gradient_predivide_factor)``) will still be applied, but the post-allreduce gradient averaging (``grads.mul_(gradient_predivide_factor/world size)``) will be omitted.
+
+    """
+
+    def __init__(self,
+                 module,
+                 message_size=10000000,
+                 delay_allreduce=False,
+                 shared_param=None,
+                 allreduce_trigger_params=None,
+                 retain_allreduce_buffers=False,
+                 allreduce_always_fp32=False,
+                 num_allreduce_streams=1,
+                 allreduce_communicators=None,
+                 gradient_average=True,
+                 gradient_predivide_factor=1.0,
+                 gradient_average_split_factor=None,
+                 prof=False):
+        super(DistributedDataParallel, self).__init__()
+
+        # Backward/forward compatibility around
+        # https://github.com/pytorch/pytorch/commit/540ef9b1fc5506369a48491af8a285a686689b36 and
+        # https://github.com/pytorch/pytorch/commit/044d00516ccd6572c0d6ab6d54587155b02a3b86
+        if hasattr(dist, "get_backend"):
+            self._backend = dist.get_backend()
+            if hasattr(dist, "DistBackend"):
+                self.backend_enum_holder = dist.DistBackend
+            else:
+                self.backend_enum_holder = dist.Backend
+        else:
+            self._backend = dist._backend
+            self.backend_enum_holder = dist.dist_backend
+
+        self.warn_on_half = True if self._backend == self.backend_enum_holder.GLOO else False
+
+        self.prof = prof
+
+        self.allreduce_different_streams = (num_allreduce_streams > 1)
+        self.num_allreduce_streams = num_allreduce_streams
+        self.allreduce_communicators = allreduce_communicators
+        if self.allreduce_communicators:
+            assert len(allreduce_communicators[0]) == num_allreduce_streams
+            assert len(allreduce_communicators[0]) == len(allreduce_communicators[1])
+            assert self.allreduce_different_streams
+
+        if self.allreduce_different_streams and delay_allreduce:
+            raise ValueError("self.allreduce_different_streams may only be used if delay_allreduce=False.")
+
+        if shared_param is not None:
+            raise ValueError("shared_param is no longer supported as an option.  It was misleadingly named from the start.  It turns out overlapping communication with computation should work fine with shared parameters.  If you still wish to delay communication to the end of the backward pass, use delay_allreduce=True|False instead.")
+
+        self.world_size = float(dist.get_world_size())
+
+        self.retain_allreduce_buffers = retain_allreduce_buffers
+        self.allreduce_always_fp32 = allreduce_always_fp32
+        self.gradient_average = gradient_average
+        self.gradient_predivide_factor = gradient_predivide_factor
+
+        self.custom_allreduce_triggers = False
+        if allreduce_trigger_params is not None:
+            if delay_allreduce:
+                raise ValueError("Setting allreduce_trigger_params is only valid if delay_allreduce=False.")
+            self.custom_allreduce_triggers = True
+            self.allreduce_trigger_params = set([id(param) for param in allreduce_trigger_params])
+
+        self.delay_allreduce = delay_allreduce
+        self.message_size = message_size
+
+        self.main_stream = torch.cuda.current_stream()
+
+        self.bucket_streams = []
+        self.bucket_events = []
+
+        self.module = module
+
+        self._disable_allreduce = False
+
+        if self._backend == self.backend_enum_holder.NCCL:
+            for param in self.module.parameters():
+                assert param.is_cuda, "NCCL backend only supports model parameters to be on GPU."
+
+        self.active_params = []
+
+        self.param_type_to_tmp_i = {"torch.cuda.HalfTensor" : 0,
+                                    "torch.cuda.FloatTensor" : 1,
+                                    "torch.cuda.DoubleTensor" : 2}
+
+        if multi_tensor_applier.available:
+            # TODO:  I really need to centralize the C++ backed imports
+            import amp_C
+            self.multi_tensor_scale = amp_C.multi_tensor_scale
+            self._overflow_buf = torch.cuda.IntTensor([0])
+
+        self.create_hooks()
+
+        flat_dist_call([param.data for param in self.module.parameters()], dist.broadcast, (0,) )
+
+
+    def __setstate__(self, state):
+        super(DistributedDataParallel, self).__setstate__(state)
+        if self.allreduce_different_streams and delay_allreduce:
+            raise ValueError("self.allreduce_different_streams may only be used if delay_allreduce=False.")
+
+        if self.delay_allreduce:
+            self.needs_refresh = True
+
+        self.bucket_streams = []
+        self.bucket_events = []
+
+
+    def __getstate__(self):
+        attrs = copy.copy(self.__dict__)
+        if self._backend != self.backend_enum_holder.NCCL:
+            del attrs['self.bucket_streams']
+            del attrs['self.bucket_events']
+            return attrs
+
+    def enable_allreduce(self):
+        self._disable_allreduce = False
+
+    def disable_allreduce(self):
+        self._disable_allreduce = True
+
+    # Broadcast rank 0's bucket structure across all processes, and have all processes
+    # regenerate their bucket structures to match.
+    def sync_bucket_structure(self):
+        # Append leftover buckets
+        for tmp_bucket in self.tmp_buckets:
+            if len(tmp_bucket) > 0:
+                self.active_i_buckets.append(tmp_bucket)
+
+        self.num_buckets = len(self.active_i_buckets)
+        self.bucket_sizes = [len(bucket) for bucket in self.active_i_buckets]
+
+        info_tensor = torch.cuda.IntTensor([self.num_buckets] +
+                                           self.bucket_sizes +
+                                           list(chain(*self.active_i_buckets)))
+
+        dist.broadcast(info_tensor, 0)
+
+        info = [int(entry) for entry in info_tensor]
+
+        self.num_buckets = info[0]
+        self.bucket_sizes = info[1:self.num_buckets + 1]
+        self.buckets = [[None for _ in range(self.bucket_sizes[i])]
+                        for i in range(self.num_buckets)]
+        # Technically, active_i_buckets' work is done.  But the information is still useful to
+        # keep around.  Therefore, refresh active_i_buckets based on rank 0 as well.
+        self.active_i_buckets = [[None for _ in range(self.bucket_sizes[i])]
+                                 for i in range(self.num_buckets)]
+
+        flattened_buckets = info[self.num_buckets + 1:]
+        flat_i = 0
+        for bucket_idx in range(self.num_buckets):
+            for bucket_loc in range(self.bucket_sizes[bucket_idx]):
+                param_i = flattened_buckets[flat_i]
+                self.active_i_buckets[bucket_idx][bucket_loc] = param_i
+                self.param_id_to_bucket[id(self.active_params[param_i])] = (bucket_idx, bucket_loc)
+                flat_i += 1
+
+
+    def create_hooks(self):
+        # Fallback hook that's only called at the end of backward.
+        # Used if you deliberately want to delay allreduces to the end, or to refresh the
+        # bucket structure that will be used to overlap communication with computation in later
+        # iterations.
+        def allreduce_params():
+            # Bucket record refresh
+            if not self.delay_allreduce:
+                if self.needs_refresh:
+                    self.sync_bucket_structure()
+
+                    self.needs_refresh = False
+
+            self.allreduce_fallback()
+
+
+        def overlapping_backward_epilogue():
+            for stream, event in zip(self.bucket_streams, self.bucket_events):
+                stream.record_event(event)
+                torch.cuda.current_stream().wait_event(event)
+
+            # Sanity checks that all the buckets were kicked off
+            if self.next_bucket != self.num_buckets:
+                raise RuntimeError("In epilogue, next_bucket ({}) != num_buckets ({}).  ".format(
+                                   self.next_bucket, self.num_buckets),
+                                   "This probably indicates some buckets were not allreduced.")
+
+            for actual, expected in zip(self.buckets_ready_size, self.bucket_sizes):
+                if actual != expected:
+                    raise RuntimeError("Some param buckets were not allreduced.")
+
+
+        self.grad_accs = []
+        for param in self.module.parameters():
+            if param.requires_grad:
+                def wrapper(param):
+                    param_tmp = param.expand_as(param)
+                    grad_acc = param_tmp.grad_fn.next_functions[0][0]
+
+                    def allreduce_hook(*unused):
+                        if self.prof:
+                            torch.cuda.nvtx.range_push("allreduce_hook")
+
+                        if not self._disable_allreduce:
+                            if self.delay_allreduce or self.needs_refresh:
+                                # TODO:  How do we want to handle multiple backward passes between
+                                # each forward, e.g., backward passes with retain_graph=True?
+                                # needs_refresh and callback_queued are both vulnerable states.
+                                if not self.delay_allreduce and self.needs_refresh:
+                                    # Use the backward pass to build the bucket structure on the fly.
+                                    active_i = self.param_id_to_active_i[id(param)]
+
+                                    # Float, half, and double tensors are grouped into buckets separately.
+                                    current_type = self.param_type_to_tmp_i[param.type()]
+
+                                    self.tmp_buckets[current_type].append(active_i)
+
+                                    ship_tmp_bucket = False
+                                    if self.custom_allreduce_triggers:
+                                        if id(param) in self.allreduce_trigger_params:
+                                            ship_tmp_bucket = True
+                                    else:
+                                        self.tmp_numels[current_type] += param.numel()
+                                        if self.tmp_numels[current_type] >= self.message_size:
+                                            ship_tmp_bucket = True
+
+                                    # To consider:  If custom_allreduce_triggers are in use, ship all
+                                    # tmp_buckets, not just tmp_buckets[current_type].
+                                    if ship_tmp_bucket:
+                                        self.active_i_buckets.append(self.tmp_buckets[current_type])
+                                        self.tmp_buckets[current_type] = []
+                                        self.tmp_numels[current_type] = 0
+
+                                if not self.callback_queued:
+                                    Variable._execution_engine.queue_callback(allreduce_params)
+                                    self.callback_queued = True
+                            else:
+                                if not self.callback_queued:
+                                    Variable._execution_engine.queue_callback(overlapping_backward_epilogue)
+                                    self.callback_queued = True
+
+                                self.comm_ready_buckets(param)
+
+                        if self.prof:
+                            torch.cuda.nvtx.range_pop()
+
+                    grad_acc.register_hook(allreduce_hook)
+                    self.grad_accs.append(grad_acc)
+
+                wrapper(param)
+
+
+    def _stream_this_bucket(self, bucket_idx):
+        if self.allreduce_different_streams:
+            return self.bucket_streams[bucket_idx%self.num_allreduce_streams]
+        else:
+            return self.bucket_streams[0]
+
+
+    def _event_this_bucket(self, bucket_idx):
+        if self.allreduce_different_streams:
+            return self.bucket_events[bucket_idx%self.num_allreduce_streams]
+        else:
+            return self.bucket_events[0]
+
+
+    def allreduce_bucket(self, bucket, bucket_idx, force_default_stream):
+        tensor = flatten(bucket)
+
+        if force_default_stream:
+            bucket_stream = self.main_stream
+        else:
+            bucket_stream = self._stream_this_bucket(bucket_idx)
+            bucket_event = self._event_this_bucket(bucket_idx)
+            torch.cuda.current_stream().record_event(bucket_event)
+            bucket_stream.wait_event(bucket_event)
+
+        with torch.cuda.stream(bucket_stream):
+            # self.main_stream.wait_stream(torch.cuda.current_stream())
+            # torch.cuda.synchronize()
+
+            tensor_to_allreduce = tensor
+
+            if self.allreduce_always_fp32:
+                tensor_to_allreduce = tensor.float()
+
+            if self.gradient_predivide_factor != 1.0:
+                tensor_to_allreduce.mul_(1./self.gradient_predivide_factor)
+
+            if self.allreduce_different_streams and not force_default_stream:
+                dist.all_reduce(tensor_to_allreduce, group=self.bucket_pgs[bucket_idx%self.num_allreduce_streams])
+            else:
+                dist.all_reduce(tensor_to_allreduce)
+
+            if self.gradient_average:
+                tensor_to_allreduce.mul_(self.gradient_predivide_factor/self.world_size)
+
+            if self.allreduce_always_fp32 and tensor is not tensor_to_allreduce:
+                tensor.copy_(tensor_to_allreduce)
+
+            if not self.retain_allreduce_buffers:
+                if multi_tensor_applier.available:
+                    multi_tensor_applier(
+                        self.multi_tensor_scale,
+                        self._overflow_buf,
+                        [unflatten(tensor, bucket), bucket],
+                        1.0)
+                else:
+                    for buf, synced in zip(bucket, unflatten(tensor, bucket)):
+                        buf.copy_(synced)
+
+            # I think we actually do need this here.  After allreduce_bucket returns, tensor will
+            # eventually go out of scope and die, at which point it could otherwise be freed for
+            # further reuse by the main stream while the allreduce/div/unflatten are underway in bucket_stream.
+            tensor.record_stream(bucket_stream)
+
+        return tensor
+
+
+    def allreduce_maybe_retain(self, bucket, bucket_idx, force_default_stream=False):
+        allreduced = self.allreduce_bucket(bucket, bucket_idx, force_default_stream)
+        if self.retain_allreduce_buffers:
+            if self.allreduce_buffers[bucket_idx] is not None:
+                raise RuntimeError("The backward pass is attempting to replace an already-filled "
+                                   "allreduce buffer.  This is almost certainly an error.")
+            self.allreduce_buffers[bucket_idx] = allreduced
+            for view, grad in zip(unflatten(allreduced, bucket), bucket):
+                grad.data = view
+            # for buf, synced in zip(bucket, unflatten(allreduced, bucket)):
+            #     buf.copy_(synced)
+
+
+    def allreduce_fallback(self):
+        for stream, event in zip(self.bucket_streams, self.bucket_events):
+            stream.record_event(event)
+            torch.cuda.current_stream().wait_event(event)
+
+        if self.retain_allreduce_buffers:
+            grads = [param.grad for param in self.module.parameters() if param.grad is not None]
+        else:
+            grads = [param.grad.data for param in self.module.parameters() if param.grad is not None]
+
+        split_buckets = split_half_float_double(grads)
+
+        # If retain_allreduce_buffers is True and delay_allreduce is False,
+        # this will only be done during the first backward pass, ignored by the
+        # training script, and overwritten in the next forward pass.  So it's harmless.
+        if self.retain_allreduce_buffers:
+            self.allreduce_buffers = [None for _ in range(len(split_buckets))]
+
+        for i, bucket in enumerate(split_buckets):
+            allreduced = self.allreduce_maybe_retain(bucket, i, force_default_stream=True)
+
+
+    def comm_ready_buckets(self, param):
+        # Need to do this in every hook for compatibility with Ruberry's streaming backward PR.
+        # self.reduction_stream.wait_stream(torch.cuda.current_stream())
+        if self.prof:
+            torch.cuda.nvtx.range_push("comm_ready_buckets")
+
+        bucket_idx, bucket_loc = self.param_id_to_bucket[id(param)]
+
+        if self.buckets[bucket_idx][bucket_loc] is not None:
+            raise RuntimeError("The backward pass is attempting to replace an already-filled "
+                               "bucket slot.  This is almost certainly an error.")
+
+        if self.retain_allreduce_buffers:
+            self.buckets[bucket_idx][bucket_loc] = param.grad
+        else:
+            self.buckets[bucket_idx][bucket_loc] = param.grad.data
+
+        self.buckets_ready_size[bucket_idx] += 1
+
+        if self.buckets_ready_size[bucket_idx] == self.bucket_sizes[bucket_idx]:
+            if bucket_idx == self.next_bucket:
+                self.allreduce_maybe_retain(self.buckets[bucket_idx], bucket_idx)
+
+                self.next_bucket += 1
+
+                # Reversing upstream's logic here, because we constructed our buckets based on
+                # the order things were received during backward.
+                if len(self.ready_buckets_not_reduced) > 0:
+                    sorted_todo = sorted(self.ready_buckets_not_reduced)
+                    for i in sorted_todo:
+                        # Nothing can be reduced now
+                        if i > self.next_bucket:
+                            break
+                        elif i == self.next_bucket:
+                            self.allreduce_maybe_retain(self.buckets[i], i)
+                            self.ready_buckets_not_reduced.remove(i)
+                            self.next_bucket += 1
+                        else:
+                            raise ValueError("i should always be >= next_bucket")
+            else:
+                self.ready_buckets_not_reduced.add(bucket_idx)
+
+        if self.prof:
+            torch.cuda.nvtx.range_pop()
+
+
+    def forward(self, *inputs, **kwargs):
+        result = self.module(*inputs, **kwargs)
+
+        if self.prof:
+            torch.cuda.nvtx.range_push("forward pass DDP logic")
+
+        if not self._disable_allreduce:
+            if not self.delay_allreduce:
+                param_list = [param for param in self.module.parameters() if param.requires_grad]
+
+                # Conditions under which to refresh self.record
+                # Forward has the authority to set needs_refresh to True, but only allreduce_params
+                # in backward has the authority to set needs_refresh to False.
+                # Parentheses are not necessary for correct order of operations, but make the intent clearer.
+                if ((not self.active_params) or
+                    (len(param_list) != len(self.active_params)) or
+                    any([param1 is not param2 for param1, param2 in zip(param_list, self.active_params)])):
+                    self.needs_refresh = True
+
+                if self.needs_refresh:
+                    self.active_i_buckets = []
+                    self.buckets = []
+                    self.tmp_buckets = [[], [], []] # [running half, float, double buckets]
+                    self.tmp_numels = [0, 0, 0]
+                    self.bucket_sizes = []
+                    self.param_id_to_active_i = {id(param) : i for i, param in enumerate(param_list)}
+                    self.param_id_to_bucket = {}
+                    self.bucket_pgs = []
+                    self.bucket_streams = []
+                    self.bucket_events = []
+                else:
+                    # self.buckets = [[None for _ in range(self.bucket_sizes[i])]
+                    #                 for i in range(self.num_buckets)]
+                    if not self.buckets:
+                        self.buckets = [[None for _ in range(self.bucket_sizes[i])]
+                                        for i in range(self.num_buckets)]
+                    else:
+                        assert len(self.buckets) == self.num_buckets, "len(buckets) = {}, expected {}".format(
+                            len(self.buckets), self.num_buckets)
+                        for b, bucket in enumerate(self.buckets):
+                            assert len(bucket) == self.bucket_sizes[b], "len(buckets[{}]) = {}, expected {})".format(
+                                b, len(buckets[b]), self.bucket_sizes[b])
+                            for i in range(len(bucket)):
+                                bucket[i] = None
+
+                    if self.allreduce_communicators:
+                        self.bucket_pgs = self.allreduce_communicators[0]
+                        self.bucket_streams = self.allreduce_communicators[1]
+                        self.bucket_events = [torch.cuda.Event(enable_timing=False,
+                                            blocking=False) for _ in range(self.num_allreduce_streams)]
+                    else:
+                        if self.allreduce_different_streams:
+                            if not self.bucket_pgs:
+                                self.bucket_pgs = [dist.new_group() for _ in range(self.num_allreduce_streams)]
+                                for i, bg in enumerate(self.bucket_pgs):
+                                    print("rank {} created group {} with backend {}".format(
+                                          dist.get_rank(), i, dist.get_backend(bg)))
+                        if self.allreduce_different_streams:
+                            if not self.bucket_streams:
+                                self.bucket_streams = [torch.cuda.Stream() for _ in range(self.num_allreduce_streams)]
+                                self.bucket_events = [torch.cuda.Event(enable_timing=False,
+                                                      blocking=False) for _ in range(self.num_allreduce_streams)]
+                        else:
+                            if not self.bucket_streams:
+                                self.bucket_streams = [torch.cuda.Stream()]
+                                self.bucket_events = [torch.cuda.Event(enable_timing=False, blocking=False)]
+
+                    self.buckets_ready_size = [0 for i in range(self.num_buckets)]
+                    if(self.retain_allreduce_buffers):
+                        self.allreduce_buffers = [None for _ in range(self.num_buckets)]
+                    self.next_bucket = 0
+                    self.ready_buckets_not_reduced = set()
+
+                self.active_params = param_list
+
+            self.callback_queued = False
+
+        if self.prof:
+            torch.cuda.nvtx.range_pop()
+
+        return result
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/multiproc.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/multiproc.py
new file mode 100644
index 0000000000..ff743df20b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/multiproc.py
@@ -0,0 +1,35 @@
+import torch
+import sys
+import subprocess
+
+def docstring_hack():
+    """
+    Multiproc file which will launch a set of processes locally for multi-gpu
+    usage: python -m apex.parallel.multiproc main.py ...
+    """
+    pass
+
+argslist = list(sys.argv)[1:]
+world_size = torch.cuda.device_count()
+
+if '--world-size' in argslist:
+    world_size = int(argslist[argslist.index('--world-size')+1])
+else:
+    argslist.append('--world-size')
+    argslist.append(str(world_size))
+
+workers = []
+
+for i in range(world_size):
+    if '--rank' in argslist:
+        argslist[argslist.index('--rank')+1] = str(i)
+    else:
+        argslist.append('--rank')
+        argslist.append(str(i))
+    stdout = None if i == 0 else open("GPU_"+str(i)+".log", "w")
+    print(argslist)
+    p = subprocess.Popen([str(sys.executable)]+argslist, stdout=stdout)
+    workers.append(p)
+
+for p in workers:
+    p.wait()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/optimized_sync_batchnorm.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/optimized_sync_batchnorm.py
new file mode 100644
index 0000000000..65cf5eabf6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/optimized_sync_batchnorm.py
@@ -0,0 +1,85 @@
+import torch
+from torch.nn.modules.batchnorm import _BatchNorm
+from torch.nn import functional as F
+
+import syncbn
+from .optimized_sync_batchnorm_kernel import SyncBatchnormFunction
+
+
+class SyncBatchNorm(_BatchNorm):
+    """
+    synchronized batch normalization module extented from `torch.nn.BatchNormNd`
+    with the added stats reduction across multiple processes.
+    :class:`apex.parallel.SyncBatchNorm` is designed to work with
+    `DistributedDataParallel`.
+
+    When running in training mode, the layer reduces stats across all processes
+    to increase the effective batchsize for normalization layer. This is useful
+    in applications where batch size is small on a given process that would
+    diminish converged accuracy of the model. The model uses collective
+    communication package from `torch.distributed`.
+
+    When running in evaluation mode, the layer falls back to
+    `torch.nn.functional.batch_norm`
+
+    Args:
+        num_features: :math:`C` from an expected input of size
+            :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average
+            (i.e. simple average). Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters. Default: ``True``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics and always uses batch
+            statistics in both training and eval modes. Default: ``True``
+        process_group: pass in a process group within which the stats of the
+            mini-batch is being synchronized. ``None`` for using default process
+            group
+        channel_last: a boolean value that when set to ``True``, this module
+            take the last dimension of the input tensor to be the channel
+            dimension. Default: False
+
+    Examples::
+        >>> # channel first tensor
+        >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()
+        >>> inp = torch.randn(10, 100, 14, 14).cuda()
+        >>> out = sbn(inp)
+        >>> inp = torch.randn(3, 100, 20).cuda()
+        >>> out = sbn(inp)
+        >>> # channel last tensor
+        >>> sbn = apex.parallel.SyncBatchNorm(100, channel_last=True).cuda()
+        >>> inp = torch.randn(10, 14, 14, 100).cuda()
+    """
+
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False, fuse_relu=False):
+        super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)
+        self.process_group = process_group
+        self.channel_last = channel_last
+        self.fuse_relu = fuse_relu
+
+    def _specify_process_group(self, process_group):
+        self.process_group = process_group
+
+    def _specify_channel_last(self, channel_last):
+        self.channel_last = channel_last
+
+    def forward(self, input, z = None):
+        # if input.dim() == 2, we switch to channel_last for efficient memory accessing
+        channel_last = self.channel_last if input.dim() != 2 else True
+
+        if not self.training and self.track_running_stats and not channel_last and not self.fuse_relu and z == None:
+            # fall back to pytorch implementation for inference
+            return F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)
+        else:
+            exponential_average_factor = 0.0
+            if self.training and self.track_running_stats:
+                self.num_batches_tracked += 1
+                if self.momentum is None:
+                    exponential_average_factor = 1.0 / float(self.num_batches_tracked)
+                else:
+                    exponential_average_factor = self.momentum
+            return SyncBatchnormFunction.apply(input, z, self.weight, self.bias, self.running_mean, self.running_var, self.eps, self.training or not self.track_running_stats, exponential_average_factor, self.process_group, channel_last, self.fuse_relu)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/optimized_sync_batchnorm_kernel.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/optimized_sync_batchnorm_kernel.py
new file mode 100644
index 0000000000..6168471496
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/optimized_sync_batchnorm_kernel.py
@@ -0,0 +1,119 @@
+import torch
+from torch.autograd.function import Function
+
+import syncbn
+from apex.parallel import ReduceOp
+
+class SyncBatchnormFunction(Function):
+
+    @staticmethod
+    def forward(ctx, input, z, weight, bias, running_mean, running_variance, eps, track_running_stats = True, momentum = 1.0, process_group = None, channel_last = False, fuse_relu = False):
+        input = input.contiguous()
+        world_size = 0
+
+        mean = None
+        var_biased = None
+        inv_std = None
+        var = None
+        out = None
+        count = None
+        if track_running_stats:
+            if channel_last:
+                count = int(input.numel()/input.size(-1))
+                mean, var_biased = syncbn.welford_mean_var_c_last(input)
+                num_channels = input.size(-1)
+            else:
+                count = int(input.numel()/input.size(1))
+                mean, var_biased = syncbn.welford_mean_var(input)
+                num_channels = input.size(1)
+
+            if torch.distributed.is_initialized():
+                if not process_group:
+                    process_group = torch.distributed.group.WORLD
+                device = mean.device
+                world_size = torch.distributed.get_world_size(process_group)
+
+                count_t = torch.empty(1, dtype=mean.dtype, device=mean.device).fill_(count)
+                combined = torch.cat([mean.view(-1), var_biased.view(-1), count_t], dim=0)
+                combined_list = [torch.empty_like(combined) for k in range(world_size)]
+                torch.distributed.all_gather(combined_list, combined, process_group)
+                combined = torch.stack(combined_list, dim=0)
+                mean_all, invstd_all, count_all = torch.split(combined, num_channels, dim=1)
+                count_all = count_all.view(-1)
+                mean, var, inv_std = syncbn.welford_parallel(mean_all, invstd_all, count_all.to(torch.int32), eps)
+            else:
+                device = mean.device
+                count_all = torch.cuda.IntTensor([count], device=device)
+                inv_std = 1.0 / torch.sqrt(var_biased + eps)
+                var = var_biased * (count) / (count-1)
+
+            if count == 1 and world_size < 2:
+                raise ValueError('Expected more than 1 value per channel when training, got input size{}'.format(input.size()))
+
+            r_m_inc = mean if running_mean.dtype != torch.float16 else mean.half()
+            r_v_inc = var if running_variance.dtype != torch.float16 else var.half()
+            running_mean.data = running_mean.data * (1-momentum) + momentum*r_m_inc
+            running_variance.data = running_variance.data * (1-momentum) + momentum*r_v_inc
+        else:
+            mean = running_mean.data
+            inv_std = 1.0 / torch.sqrt(running_variance.data + eps)
+
+        ctx.save_for_backward(input, weight, mean, inv_std, z, bias, count_all.to(torch.int32))
+        ctx.process_group = process_group
+        ctx.channel_last = channel_last
+        ctx.world_size = world_size
+        ctx.fuse_relu = fuse_relu
+
+        if channel_last:
+            out = syncbn.batchnorm_forward_c_last(input, z, mean, inv_std, weight, bias, fuse_relu)
+        else:
+            out = syncbn.batchnorm_forward(input, mean, inv_std, weight, bias)
+
+        return out
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        grad_output = grad_output.contiguous()
+        # mini batch mean & var are calculated by forward path.
+        # mu = 1./N*np.sum(h, axis = 0)
+        # var = 1./N*np.sum((h-mu)**2, axis = 0)
+        saved_input, weight, mean, inv_std, z, bias, count = ctx.saved_tensors
+        process_group = ctx.process_group
+        channel_last = ctx.channel_last
+        world_size = ctx.world_size
+        fuse_relu = ctx.fuse_relu
+        grad_input = grad_z = grad_weight = grad_bias = None
+
+        if fuse_relu:
+            grad_output = syncbn.relu_bw_c_last(grad_output, saved_input, z, mean, inv_std, weight, bias)
+        if isinstance(z, torch.Tensor) and ctx.needs_input_grad[1]:
+            grad_z = grad_output.clone()
+
+        # TODO: update kernel to not pre_divide by item_num
+        if channel_last:
+            sum_dy, sum_dy_xmu, grad_weight, grad_bias = syncbn.reduce_bn_c_last(grad_output, saved_input, mean, inv_std, weight)
+        else:
+            sum_dy, sum_dy_xmu, grad_weight, grad_bias = syncbn.reduce_bn(grad_output, saved_input, mean, inv_std, weight)
+
+        # calculate grad_input
+        if ctx.needs_input_grad[0]:
+
+            if torch.distributed.is_initialized():
+                num_channels = sum_dy.shape[0]
+                combined = torch.cat([sum_dy, sum_dy_xmu], dim=0)
+                torch.distributed.all_reduce(
+                    combined, torch.distributed.ReduceOp.SUM, process_group, async_op=False)
+                sum_dy, sum_dy_xmu = torch.split(combined, num_channels)
+
+            if channel_last:
+                grad_input = syncbn.batchnorm_backward_c_last(grad_output, saved_input, mean, inv_std, weight, sum_dy, sum_dy_xmu, count)
+            else:
+                grad_input = syncbn.batchnorm_backward(grad_output, saved_input, mean, inv_std, weight, sum_dy, sum_dy_xmu, count)
+
+        if weight is None or not ctx.needs_input_grad[2]:
+            grad_weight = None
+
+        if weight is None or not ctx.needs_input_grad[3]:
+            grad_bias = None
+
+        return grad_input, grad_z, grad_weight, grad_bias, None, None, None, None, None, None, None, None
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/sync_batchnorm.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/sync_batchnorm.py
new file mode 100644
index 0000000000..1fcfe4356a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/sync_batchnorm.py
@@ -0,0 +1,134 @@
+import torch
+from torch.nn.modules.batchnorm import _BatchNorm
+from torch.nn import functional as F
+
+from .sync_batchnorm_kernel import SyncBatchnormFunction
+from apex.parallel import ReduceOp
+
+
+class SyncBatchNorm(_BatchNorm):
+    """
+    synchronized batch normalization module extented from ``torch.nn.BatchNormNd``
+    with the added stats reduction across multiple processes.
+    :class:`apex.parallel.SyncBatchNorm` is designed to work with
+    ``DistributedDataParallel``.
+
+    When running in training mode, the layer reduces stats across all processes
+    to increase the effective batchsize for normalization layer. This is useful
+    in applications where batch size is small on a given process that would
+    diminish converged accuracy of the model. The model uses collective
+    communication package from ``torch.distributed``.
+
+    When running in evaluation mode, the layer falls back to
+    ``torch.nn.functional.batch_norm``.
+
+    Args:
+        num_features: :math:`C` from an expected input of size
+            :math:`(N, C, L)` or :math:`L` from input of size :math:`(N, L)`
+        eps: a value added to the denominator for numerical stability.
+            Default: 1e-5
+        momentum: the value used for the running_mean and running_var
+            computation. Can be set to ``None`` for cumulative moving average
+            (i.e. simple average). Default: 0.1
+        affine: a boolean value that when set to ``True``, this module has
+            learnable affine parameters. Default: ``True``
+        track_running_stats: a boolean value that when set to ``True``, this
+            module tracks the running mean and variance, and when set to ``False``,
+            this module does not track such statistics and always uses batch
+            statistics in both training and eval modes. Default: ``True``
+
+    Example::
+
+        >>> sbn = apex.parallel.SyncBatchNorm(100).cuda()
+        >>> inp = torch.randn(10, 100, 14, 14).cuda()
+        >>> out = sbn(inp)
+        >>> inp = torch.randn(3, 100, 20).cuda()
+        >>> out = sbn(inp)
+    """
+
+    warned = False
+
+    def __init__(self, num_features, eps=1e-5, momentum=0.1, affine=True, track_running_stats=True, process_group=None, channel_last=False):
+        if channel_last == True:
+            raise AttributeError("channel_last is not supported by primitive SyncBatchNorm implementation. Try install apex with `--cuda_ext` if channel_last is desired.")
+
+        if not SyncBatchNorm.warned:
+            if hasattr(self, "syncbn_import_error"):
+                print("Warning:  using Python fallback for SyncBatchNorm, possibly because apex was installed without --cuda_ext.  The exception raised when attempting to import the cuda backend was: ", self.syncbn_import_error)
+            else:
+                print("Warning:  using Python fallback for SyncBatchNorm")
+            SyncBatchNorm.warned = True
+
+        super(SyncBatchNorm, self).__init__(num_features, eps=eps, momentum=momentum, affine=affine, track_running_stats=track_running_stats)
+        self.process_group = process_group
+
+    def _specify_process_group(self, process_group):
+        self.process_group = process_group
+
+    def forward(self, input):
+        torch.cuda.nvtx.range_push("sync_bn_fw_with_mean_var")
+        mean = None
+        var = None
+        cast = None
+        out = None
+
+        # casting to handle mismatch input type to layer type
+        if self.running_mean is not None:
+            if self.running_mean.dtype != input.dtype:
+                input = input.to(self.running_mean.dtype)
+                cast = input.dtype
+        elif self.weight is not None:
+            if self.weight.dtype != input.dtype:
+                input = input.to(self.weight.dtype)
+                cast = input.dtype
+
+        if not self.training and self.track_running_stats:
+            # fall back to pytorch implementation for inference
+            torch.cuda.nvtx.range_pop()
+            out = F.batch_norm(input, self.running_mean, self.running_var, self.weight, self.bias, False, 0.0, self.eps)
+        else:
+            process_group = self.process_group
+            world_size = 1
+            if not self.process_group:
+                process_group = torch.distributed.group.WORLD
+            self.num_batches_tracked += 1
+            with torch.no_grad():
+                channel_first_input = input.transpose(0, 1).contiguous()
+                squashed_input_tensor_view = channel_first_input.view(
+                    channel_first_input.size(0), -1)
+                # total number of data points for each variance entry. Used to calculate unbiased variance estimate
+                m = None
+                local_m = float(squashed_input_tensor_view.size()[1])
+                local_mean = torch.mean(squashed_input_tensor_view, 1)
+                local_sqr_mean = torch.pow(
+                    squashed_input_tensor_view, 2).mean(1)
+                if torch.distributed.is_initialized():
+                    world_size = torch.distributed.get_world_size(process_group)
+                    torch.distributed.all_reduce(
+                        local_mean, ReduceOp.SUM, process_group)
+                    mean = local_mean / world_size
+                    torch.distributed.all_reduce(
+                        local_sqr_mean, ReduceOp.SUM, process_group)
+                    sqr_mean = local_sqr_mean / world_size
+                    m = local_m * world_size
+                else:
+                    m = local_m
+                    mean = local_mean
+                    sqr_mean = local_sqr_mean
+                # var(x) = E (( x - mean_x ) ** 2)
+                #        = 1 / N * sum ( x - mean_x ) ** 2
+                #        = 1 / N * sum (x**2) - mean_x**2
+                var = sqr_mean - mean.pow(2)
+
+                if self.running_mean is not None:
+                    self.running_mean = self.momentum * mean + \
+                        (1 - self.momentum) * self.running_mean
+                if self.running_var is not None:
+                    # as noted by the paper, we used unbiased variance estimate of the mini-batch
+                    # Var[x] = m / (m-1) * Eb (sample_variance)
+                    self.running_var = m / \
+                        (m-1) * self.momentum * var + \
+                        (1 - self.momentum) * self.running_var
+            torch.cuda.nvtx.range_pop()
+            out = SyncBatchnormFunction.apply(input, self.weight, self.bias, mean, var, self.eps, process_group, world_size)
+        return out.to(cast)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/sync_batchnorm_kernel.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/sync_batchnorm_kernel.py
new file mode 100644
index 0000000000..e407a63da3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/parallel/sync_batchnorm_kernel.py
@@ -0,0 +1,87 @@
+import torch
+from torch.autograd.function import Function
+
+from apex.parallel import ReduceOp
+
+
+class SyncBatchnormFunction(Function):
+
+    @staticmethod
+    def forward(ctx, input, weight, bias, running_mean, running_variance, eps, process_group, world_size):
+        torch.cuda.nvtx.range_push("sync_BN_fw")
+        # transpose it to channel last to support broadcasting for input with different rank
+        c_last_input = input.transpose(1, -1).contiguous().clone()
+
+        ctx.save_for_backward(c_last_input, weight, bias,
+                              running_mean, running_variance)
+        ctx.eps = eps
+        ctx.process_group = process_group
+        ctx.world_size = world_size
+
+        c_last_input = (c_last_input - running_mean) / \
+            torch.sqrt(running_variance + eps)
+
+        if weight is not None:
+            c_last_input = c_last_input * weight
+        if bias is not None:
+            c_last_input = c_last_input + bias
+
+        torch.cuda.nvtx.range_pop()
+        return c_last_input.transpose(1, -1).contiguous().clone()
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        torch.cuda.nvtx.range_push("sync_BN_bw")
+        # mini batch mean & var are calculated by forward path.
+        # mu = 1./N*np.sum(h, axis = 0)
+        # var = 1./N*np.sum((h-mu)**2, axis = 0)
+        c_last_input, weight, bias, running_mean, running_variance = ctx.saved_tensors
+
+        eps = ctx.eps
+        process_group = ctx.process_group
+        world_size = ctx.world_size
+        grad_input = grad_weight = grad_bias = None
+        num_features = running_mean.size()[0]
+
+        # transpose it to channel last to support broadcasting for input with different rank
+        torch.cuda.nvtx.range_push("carilli field")
+        c_last_grad = grad_output.transpose(1, -1).contiguous()
+        # squash non-channel dimension so we can easily calculate mean
+        c_grad = c_last_grad.view(-1, num_features).contiguous()
+        torch.cuda.nvtx.range_pop()
+
+        # calculate grad_input
+        if ctx.needs_input_grad[0]:
+            # dh = gamma * (var + eps)**(-1. / 2.) * (dy - np.mean(dy, axis=0)
+            #     - (h - mu) * (var + eps)**(-1.0) * np.mean(dy * (h - mu), axis=0))
+            mean_dy = c_grad.mean(0)
+            mean_dy_xmu = (c_last_grad * (c_last_input -
+                                          running_mean)).view(-1, num_features).mean(0)
+            if torch.distributed.is_initialized():
+                torch.distributed.all_reduce(
+                    mean_dy, ReduceOp.SUM, process_group)
+                mean_dy = mean_dy / world_size
+                torch.distributed.all_reduce(
+                    mean_dy_xmu, ReduceOp.SUM, process_group)
+                mean_dy_xmu = mean_dy_xmu / world_size
+            c_last_grad_input = (c_last_grad - mean_dy - (c_last_input - running_mean) / (
+                running_variance + eps) * mean_dy_xmu) / torch.sqrt(running_variance + eps)
+            if weight is not None:
+                c_last_grad_input.mul_(weight)
+            grad_input = c_last_grad_input.transpose(1, -1).contiguous()
+
+        # calculate grad_weight
+        grad_weight = None
+        if weight is not None and ctx.needs_input_grad[1]:
+            # dgamma = np.sum((h - mu) * (var + eps)**(-1. / 2.) * dy, axis=0)
+            grad_weight = ((c_last_input - running_mean) / torch.sqrt(
+                running_variance + eps) * c_last_grad).view(-1, num_features).sum(0)
+
+        # calculate grad_bias
+        grad_bias = None
+        if bias is not None and ctx.needs_input_grad[2]:
+            # dbeta = np.sum(dy, axis=0)
+            grad_bias = c_grad.sum(0)
+
+        torch.cuda.nvtx.range_pop()
+        return grad_input, grad_weight, grad_bias, None, None, None, None, None
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/FAQs.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/FAQs.md
new file mode 100644
index 0000000000..27a0e97c56
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/FAQs.md
@@ -0,0 +1,21 @@
+1. How do I intercept the Adam optimizer in APEX ?
+
+	```python
+	from apex import pyprof
+	import fused_adam_cuda
+	pyprof.nvtx.wrap(fused_adam_cuda, 'adam')
+	```
+
+2. If you are using JIT and/or AMP, the correct initialization sequence is
+	1. Let any JIT to finish.
+	2. Initlialize pyprof `pyprof.nvtx.init()`.
+	3. Initialize AMP.
+
+3. How do I profile with `torch.distributed.launch` ?
+
+	```python
+	nvprof -f -o net%p.sql \
+		--profile-from-start off \
+		--profile-child-processes \
+		python -m torch.distributed.launch net.py
+	```
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/README.md
new file mode 100644
index 0000000000..ca10c46415
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/README.md
@@ -0,0 +1,252 @@
+## PyProf - PyTorch Profiling tool
+
+### What does this tool do?                                                                                                                                                                                                                  
+
+Analyzing the performance of deep neural networks is hard. Getting kernels out of [NvProf]([https://developer.nvidia.com/nvidia-visual-profiler](https://developer.nvidia.com/nvidia-visual-profiler)) or [NSight Compute]([https://developer.nvidia.com/nsight-compute](https://developer.nvidia.com/nsight-compute)) provides some generic kernel name and its execution time, but not detailed information regarding the following:
+
+ - Which layer launched it: e.g. the association of `ComputeOffsetsKernel` with a concrete PyTorch layer or API is not obvious.
+ - What the tensor dimensions and precision were: without knowing the tensor dimensions and precision, it's impossible to reason about whether the actual (silicon) kernel time is close to maximum performance of such a kernel on the GPU. Knowing the tensor dimensions and precision, we can figure out the FLOPs and bandwidth required by a layer, and then determine how close to maximum performance the kernel is for that operation.
+ - Forward-backward correlation: currently it's very hard to determine what the forward pass step was that resulted in the particular weight and data gradients (wgrad, dgrad), which makes it difficult to determine the tensor dimensions required by these backprop steps to assess their performance.
+ - Did the kernel use [Tensor Cores]([https://www.youtube.com/watch?v=yyR0ZoCeBO8](https://www.youtube.com/watch?v=yyR0ZoCeBO8))?
+ - Which line in the user's code resulted in launching this particular kernel (program trace)?
+
+PyProf addresses all of the issues above by:
+
+ 1. Instrumenting PyTorch operations to capture the tensor dimensions and precision using [NVTX](https://devblogs.nvidia.com/cuda-pro-tip-generate-custom-application-profile-timelines-nvtx). This information is recorded at profile capture time, e.g. using [NvProf](https://developer.nvidia.com/nvidia-visual-profiler).
+ 2. Querying the record produced by the profiler to correlate the kernel name and duration with PyTorch API/layer name, tensor dimensions, tensor precision, as well as calculating FLOPs and bandwidth for common operations. In addition, extra information from the profile is added for use by CUDA professionals, such as CUDA launch parameters (block/grid dimensions).
+
+Regarding FLOP and bandwidth implementations, these are usually quite straightforward. For example, for matrices A<sub>MxK</sub> and B<sub>KxN</sub>, the FLOP count for a matrix multiplication is 2 * M * N * K, and bandwidth is M * K + N * K + M * N. Note that these numbers are based on the algorithm, not the actual performance of the specific kernel. For more details, see NVIDIA's [Deep Learning Performance Guide](https://docs.nvidia.com/deeplearning/sdk/dl-performance-guide/index.html).
+
+Armed with such information, the user can determine various issues to help them tune the network. For instance, according to the [Tensor Core Performance Guide]([https://docs.nvidia.com/deeplearning/sdk/dl-performance-guide/index.html](https://docs.nvidia.com/deeplearning/sdk/dl-performance-guide/index.html)), the M, N and K dimensions that result in Tensor Core usage need to be divisible by 8. In fact, PyProf comes with a flag that lets the user obtain information regarding whether Tensor Cores were used by the kernel. Other useful information might include knowing that a particular kernel did not exploit much thread parallelism, as determined by the grid/block dimensions. Since many PyTorch kernels are open-source (or even custom written by the user, as in [CUDA Extensions]([https://pytorch.org/tutorials/advanced/cpp_extension.html](https://pytorch.org/tutorials/advanced/cpp_extension.html))), this provides the user with information that helps root cause performance issues and prioritize optimization work.
+
+
+### How to get started?
+
+1. Add the following lines to your PyTorch network:
+
+    ```python
+    import torch.cuda.profiler as profiler
+    from apex import pyprof
+    pyprof.nvtx.init()
+    ```
+
+    Run the training/inference loop with the [PyTorch's NVTX context manager](https://pytorch.org/docs/stable/_modules/torch/autograd/profiler.html#emit_nvtx)
+    `with torch.autograd.profiler.emit_nvtx()`. Optionally, you can
+    use `profiler.start()` and `profiler.stop()` to pick an iteration
+    (say after warm-up) for which you would like to capture data.
+    Here's an example:
+
+    ```python
+    iters = 500
+    iter_to_capture = 100
+
+    # Define network, loss function, optimizer etc.
+
+    # PyTorch NVTX context manager
+    with torch.autograd.profiler.emit_nvtx():
+
+        for iter in range(iters):
+
+            if iter == iter_to_capture:
+                profiler.start()
+
+            output = net(images)
+            loss = criterion(output, labels)
+            loss.backward()
+            optimizer.step()
+
+            if iter == iter_to_capture:
+                profiler.stop()
+    ```
+
+2. Run NVprof to generate a SQL (NVVP) file. This file can be opened with NVVP, as usual.
+    ```sh
+    # If you used profiler.start() and profiler.stop() in net.py
+    nvprof -f -o net.sql --profile-from-start off -- python net.py
+
+    # Profile everything
+    nvprof -f -o net.sql -- python net.py
+    ```
+
+**Note:** if you're experiencing issues with hardware counters and you get a message such as `**_ERR_NVGPUCTRPERM The user running <tool_name/application_name> does not have permission to access NVIDIA GPU Performance Counters on the target device_**`, please follow the steps described in [Hardware Counters](#hardware-counters).
+
+3. Run parser on the SQL file. The output is an ASCII file. Each line
+is a python dictionary which contains information about the kernel name,
+duration, parameters etc. This file can be used as input to other custom
+scripts as well.
+
+    ```sh
+    python -m apex.pyprof.parse net.sql > net.dict
+    ```
+
+4. Run the profiler. The input is the python dictionary created above. The tool can produce a CSV output, a columnated output (similar to `column -t` for terminal readability) and a space separated output (for post processing by AWK for instance). The tool produces 20 columns of information for every GPU kernel but you can select a subset of columns using the `-c` flag. Note that a few columns might have the value "na" implying either its a work in progress or the tool was unable to extract that information. Assuming the directory is `prof`, here are a few examples of how to use `prof.py`.
+
+    ```sh
+	# Print usage and help. Lists all available output columns.
+    python -m apex.pyprof.prof -h
+
+	# Columnated output of width 150 with some default columns.
+    python -m apex.pyprof.prof -w 150 net.dict
+
+	# CSV output.
+    python -m apex.pyprof.prof --csv net.dict
+
+	# Space seperated output.
+    python -m apex.pyprof.prof net.dict
+
+	# Columnated output of width 130 with columns index,direction,kernel name,parameters,silicon time.
+    python -m apex.pyprof.prof -w 130 -c idx,dir,kernel,params,sil net.dict
+
+	# CSV output with columns index,direction,kernel name,parameters,silicon time.
+    python -m apex.pyprof.prof --csv -c idx,dir,kernel,params,sil net.dict
+
+	# Space separated output with columns index,direction,kernel name,parameters,silicon time.
+    python -m apex.pyprof.prof -c idx,dir,kernel,params,sil net.dict
+
+	# Input redirection.
+    python -m apex.pyprof.prof < net.dict
+    ```
+
+5. Profile-guided optimization
+
+If kernels that do matrix multiplication/GEMM or convolution use half precision (fp16) data but do not use Tensor Cores (the TC column in the profile analysis output doesn't show a "1"), one can follow some basic steps to increase the likelihood that a Tensor Core-compatible kernel will be chosen. For example, for GEMMs, M, N and K should be divisible by 8, and for convolutions, the number of input and output channels shuold be divisible by 8. For more information, see detailed Tensor Core guides such as:
+- Blog Post: [Tips for Optimizing GPU Performance Using Tensor Cores](https://devblogs.nvidia.com/optimizing-gpu-performance-tensor-cores/)
+- GTC Talk: [Tensor Core Deep Learning Performance Guide](https://developer.download.nvidia.com/video/gputechconf/gtc/2019/presentation/s9926-tensor-core-performance-the-ultimate-guide.pdf)
+
+For both Tensor Core and non-Tensor Core Deep Learning performance optimization tips, see NVIDIA's [Deep Learning Performance Guide](https://docs.nvidia.com/deeplearning/sdk/dl-performance-guide/index.html).
+
+### TODOs
+1. The support for conv transpose is currently missing.
+2. PyProf currently works only with NvProf, but Nsight Compute support will be added in the future.
+
+### Example
+
+1. Run `nvprof` on the LeNet model in `examples/lenet.py`. This will output a SQL file called `net.sql`.
+
+```sh
+nvprof -f -o net.sql --profile-from-start off -- python examples/lenet.py
+```
+
+**Note**: DO NOT add --analysis-metrics since that will change which table nvprof writes the kernels to (`CUPTI_ACTIVITY_KIND_KERNEL` instead of the usual `CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL`). Support for running with metrics may be added in the future.
+
+If you don't care about a full correlation analysis and you'd just like to view the timeline with detailed NVTX annotations, you can do so, e.g. in the NVIDIA Visual Profiler (NVVP). For example, you can call `nvvp net.sql` to view the annotated timeline.
+
+2. Run the `parse.py` script on `net.sql` to extract kernel and runtime information and
+save it as `net.dict`.
+
+```sh
+python -m apex.pyprof.parse net.sql > net.dict
+```
+
+This will produce a text file, which can be parsed by any external tool, but it can also be directly read one line at a time by Python by calling `eval` on the line being read. 
+
+**Note: you do not need to process this output manually.**  Here the output is just shown as an example of modularity - you can process the raw data yourself, or let the next step enrich the information further and dump a CSV.
+
+The output of this step will look as follows. Note that the dictionary has a lot more keys than the ones shown in the example.
+
+```
+>>> with open('torchvision.resnet50.adam.64.dict') as f:
+...     for line in f:
+...         d = eval(line)
+...         print(d['kShortName'], d['op'], d['kDuration'], d['block'], d['grid'], d['device'], d['stream'], d['trace'])
+... 
+nchwToNhwc3To4Kernel ['conv2d'] 376324 (256, 1, 1) (1568, 1, 64) 0 7 ['imagenet.py:137', 'imagenet.py:129', '/opt/conda/lib/python3.6/site-packages/torchvision/models/resnet.py:195']
+generic4Channel_kernel ['conv2d'] 10720 (512, 1, 1) (19, 1, 1) 0 7 ['imagenet.py:137', 'imagenet.py:129', '/opt/conda/lib/python3.6/site-packages/torchvision/models/resnet.py:195']
+first_layer_fwd_kernel ['conv2d'] 411204 (128, 1, 1) (2, 7, 64) 0 7 ['imagenet.py:137', 'imagenet.py:129', '/opt/conda/lib/python3.6/site-packages/torchvision/models/resnet.py:195']
+nhwcToNchwKernel ['conv2d'] 342371 (256, 1, 1) (392, 2, 64) 0 7 ['imagenet.py:137', 'imagenet.py:129', '/opt/conda/lib/python3.6/site-packages/torchvision/models/resnet.py:195']
+elementwise_kernel ['__iadd__'] 2816 (128, 1, 1) (1, 1, 1) 0 7 ['imagenet.py:137', 'imagenet.py:129', '/opt/conda/lib/python3.6/site-packages/torchvision/models/resnet.py:196']
+batch_norm_collect_statistics_kernel ['batch_norm', 'batch_norm'] 929513 (512, 1, 1) (64, 1, 1) 0 7 ['imagenet.py:137', 'imagenet.py:129', '/opt/conda/lib/python3.6/site-packages/torchvision/models/resnet.py:196']
+```
+
+3. Run the `prof.py` script on `net.dict` to summarize the results into a CSV file, or to display the pretty-printed results on the screen. This step processes the raw output from step 2 to generate a nice output, but it also adds a lot of extra useful information inferred from the previous step, such as:
+- FLOPs
+- bandwidth (bytes in and out of GPU DRAM)
+- tensor core usage
+
+```sh
+python -m apex.pyprof.prof --csv net.dict > results.csv
+```
+
+You can choose which columns you'd like to display. Here's a list from calling `python -m apex.pyprof.prof -h`:
+
+```
+              idx:      Index
+              seq:      PyTorch Sequence Id
+              altseq:   PyTorch Alternate Sequence Id
+              tid:      Thread Id
+              layer:    User annotated NVTX string (can be nested)
+              trace:    Function Call Trace
+              dir:      Direction
+              sub:      Sub Sequence Id
+              mod:      Module
+              op:       Operation
+              kernel:   Kernel Name
+              params:   Parameters
+              sil:      Silicon Time (in ns)
+              tc:       Tensor Core Usage
+              device:   GPU Device Id
+              stream:   Stream Id
+              grid:     Grid Dimensions
+              block:    Block Dimensions
+              flops:    Floating point ops (FMA = 2 FLOPs)
+              bytes:    Number of bytes in and out of DRAM
+```              
+
+Let's have a look at the pretty-printed output:
+```
+python -m apex.pyprof.prof -w 100 -c kernel,op,sil,tc,flops,bytes,device,stream,block,grid torchvision.resnet50.adam.64.dict
+
+Kernel              Op              Sil(ns)    TC FLOPs        Bytes        Dev Str Block        Grid         
+elementwise_kernel  relu                381028 -      51380224    205520896   0   7 512,1,1      100352,1,1   
+volta_fp16_s884cudn conv2d              160002 1    1644167168     51388416   0   7 256,1,1      784,1,1      
+elementwise_kernel  relu                 96545 -      12845056     51380224   0   7 512,1,1      25088,1,1    
+volta_fp16_s884cudn conv2d              346083 1    6576668672    128483328   0   7 256,1,1      784,2,1      
+```
+
+Not using the pretty-print width (`-w`) option and adding `--csv` results in a CSV output instead:
+
+```
+python -m apex.pyprof.prof --csv -c kernel,mod,op,dir,sil,tc,flops,bytes,device,stream,block,grid torchvision.resnet50.adam.64.dict
+
+"Kernel","Module","Op","Direction","Sil(ns)","TC","FLOPs","Bytes","Device","Stream","Block","Grid"
+"nchwToNhwc3To4Kernel","torch.nn.functional","conv2d","fprop","376324","-","0","0","0","7","256,1,1","1568,1,64"
+"generic4Channel_kernel","torch.nn.functional","conv2d","fprop","10720","-","0","0","0","7","512,1,1","19,1,1"
+"first_layer_fwd_kernel","torch.nn.functional","conv2d","fprop","411204","-","0","0","0","7","128,1,1","2,7,64"
+"nhwcToNchwKernel","torch.nn.functional","conv2d","fprop","342371","-","0","0","0","7","256,1,1","392,2,64"
+"elementwise_kernel","Tensor","__iadd__","fprop","2816","-","1.0","8","0","7","128,1,1","1,1,1"
+"batch_norm_collect_statistics_kernel","torch.nn.functional","batch_norm","fprop","929513","-","411041792","411041792","0","7","512,1,1","64,1,1"
+"batch_norm_transform_input_kernel","torch.nn.functional","batch_norm","fprop","377539","-","411041792","411041792","0","7","512,1,1","64,64,1"
+"elementwise_kernel","torch.nn.functional","relu","fprop","381028","-","51380224","205520896","0","7","512,1,1","100352,1,1"
+"MaxPoolForward","torch.nn.functional","max_pool2d","fprop","406531","-","0","0","0","7","256,1,1","50176,1,1"
+"cudnn::gemm::computeOffsetsKernel","torch.nn.functional","conv2d","fprop","2464","-","0","0","0","7","128,1,1","25,1,1"
+```
+
+### Hardware Counters
+
+Profiling GPU workloads may require access to [hardware performance counters]([https://en.wikipedia.org/wiki/Hardware_performance_counter](https://en.wikipedia.org/wiki/Hardware_performance_counter)). Due to a [fix](https://nvidia.custhelp.com/app/answers/detail/a_id/4738) in recent NVIDIA drivers addressing [CVE‑2018‑6260](https://nvd.nist.gov/vuln/detail/CVE-2018-6260), the hardware counters are disabled by default, and require elevated privileges to be enabled again. If you're using a recent driver, you may see the following message when trying to run nvprof:
+
+```**_ERR_NVGPUCTRPERM The user running <tool_name/application_name> does not have permission to access NVIDIA GPU Performance Counters on the target device._**```
+
+For details, see [here](https://developer.nvidia.com/nvidia-development-tools-solutions-ERR_NVGPUCTRPERM-permission-issue-performance-counters).
+
+_Permanent solution_
+
+Follow the steps [here]([https://developer.nvidia.com/nvidia-development-tools-solutions-ERR_NVGPUCTRPERM-permission-issue-performance-counters](https://developer.nvidia.com/nvidia-development-tools-solutions-ERR_NVGPUCTRPERM-permission-issue-performance-counters)). The current steps for Linux are:
+```
+sudo systemctl isolate multi-user
+sudo modprobe -r nvidia_uvm nvidia_drm nvidia_modeset nvidia-vgpu-vfio nvidia
+sudo modprobe nvidia NVreg_RestrictProfilingToAdminUsers=0
+sudo systemctl isolate graphical
+```
+The above steps should result in a permanent change.
+
+_Temporary solution_
+
+When running on bare metal, you can run nvprof with `sudo`.
+
+If you're running in a Docker image, you can temporarily elevate your privileges with one of the following (oldest to newest syntax):
+<pre>
+nvidia-docker run <b>--privileged</b>
+docker run --runtime nvidia <b>--privileged</b>
+docker run --gpus all <b>--privileged<b>
+</pre>
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/__init__.py
new file mode 100644
index 0000000000..3c36146fd7
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/__init__.py
@@ -0,0 +1,3 @@
+import warnings
+
+from . import nvtx, prof
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/.gitignore b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/.gitignore
new file mode 100644
index 0000000000..b6e239e2ee
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/.gitignore
@@ -0,0 +1,4 @@
+__pycache__
+*.sql
+*.dict
+*.csv
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/README.md
new file mode 100644
index 0000000000..6af798d343
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/README.md
@@ -0,0 +1 @@
+This directory has examples of how to use `pyprof` with APEX extensions e.g. `fused_adam_cuda` and `fused_layer_norm_cuda`.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/fused_adam.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/fused_adam.py
new file mode 100644
index 0000000000..0da8b3f1ce
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/fused_adam.py
@@ -0,0 +1,20 @@
+import torch
+import fused_adam_cuda
+from apex.optimizers import FusedAdam, FP16_Optimizer
+from apex import pyprof
+
+pyprof.nvtx.init()
+pyprof.nvtx.wrap(fused_adam_cuda, 'adam')
+
+model = torch.nn.Linear(10, 20).cuda().half()
+criterion = torch.nn.CrossEntropyLoss().cuda()
+optimizer = FusedAdam(model.parameters())
+optimizer = FP16_Optimizer(optimizer)
+
+x = torch.ones(32, 10).cuda().half()
+target = torch.empty(32, dtype=torch.long).random_(20).cuda()
+y = model(x)
+loss = criterion(y, target)
+optimizer.zero_grad()
+loss.backward()
+optimizer.step()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/fused_layer_norm.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/fused_layer_norm.py
new file mode 100644
index 0000000000..2c924e4722
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/fused_layer_norm.py
@@ -0,0 +1,28 @@
+import torch
+import fused_layer_norm_cuda
+from apex.normalization import FusedLayerNorm
+from apex import pyprof
+
+pyprof.nvtx.init()
+pyprof.nvtx.wrap(fused_layer_norm_cuda, 'forward')
+pyprof.nvtx.wrap(fused_layer_norm_cuda, 'backward')
+pyprof.nvtx.wrap(fused_layer_norm_cuda, 'forward_affine')
+pyprof.nvtx.wrap(fused_layer_norm_cuda, 'backward_affine')
+
+input = torch.randn(20, 5, 10, 10).cuda()
+
+# With Learnable Parameters
+m = FusedLayerNorm(input.size()[1:]).cuda()
+output = m(input)
+
+# Without Learnable Parameters
+m = FusedLayerNorm(input.size()[1:], elementwise_affine=False).cuda()
+output = m(input)
+
+# Normalize over last two dimensions
+m = FusedLayerNorm([10, 10]).cuda()
+output = m(input)
+
+# Normalize over last dimension of size 10
+m = FusedLayerNorm(10).cuda()
+output = m(input)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/test.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/test.sh
new file mode 100644
index 0000000000..dc02f388d5
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/apex/test.sh
@@ -0,0 +1,30 @@
+#!/bin/bash
+
+set -e
+
+SCRIPT=`realpath $0`
+SCRIPTPATH=`dirname $SCRIPT`
+PYPROF="$SCRIPTPATH/../.."
+
+parse="python $PYPROF/parse/parse.py"
+prof="python $PYPROF/prof/prof.py"
+
+for f in *.py
+do
+	base=`basename $f .py`
+	sql=$base.sql
+	dict=$base.dict
+
+	#NVprof
+	echo "nvprof -fo $sql python $f"
+	nvprof -fo $sql python $f
+
+	#Parse
+	echo $parse $sql
+	$parse $sql > $dict
+
+	#Prof
+	echo $prof $dict
+	$prof -w 130 $dict
+	\rm $sql $dict
+done
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/README.md
new file mode 100644
index 0000000000..695b02b3b8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/README.md
@@ -0,0 +1 @@
+This directory has examples which show how to intercept (monkey patch) custom functions and modules with `pyprof`. No changes are required in `pyprof/parse`, however, users can add support for bytes and flops calculation for custom functions and modules in `pyprof/prof` by extending the `OperatorLayerBase` class.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/custom_function.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/custom_function.py
new file mode 100644
index 0000000000..535df2fc24
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/custom_function.py
@@ -0,0 +1,33 @@
+#!/usr/bin/env python3
+
+import torch
+import torch.cuda.profiler as profiler
+from apex import pyprof
+#Initialize pyprof
+pyprof.nvtx.init()
+
+class Foo(torch.autograd.Function):
+	@staticmethod
+	def forward(ctx, in1, in2):
+		out = in1 + in2		#This could be a custom C/C++ function.
+		return out
+
+	@staticmethod
+	def backward(ctx, grad):
+		in1_grad = grad		#This could be a custom C/C++ function.
+		in2_grad = grad		#This could be a custom C/C++ function.
+		return in1_grad, in2_grad
+
+#Hook the forward and backward functions to pyprof
+pyprof.nvtx.wrap(Foo, 'forward')
+pyprof.nvtx.wrap(Foo, 'backward')
+
+foo = Foo.apply
+
+x = torch.ones(4,4).cuda()
+y = torch.ones(4,4).cuda()
+
+with torch.autograd.profiler.emit_nvtx():
+	profiler.start()
+	z = foo(x,y)
+	profiler.stop()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/custom_module.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/custom_module.py
new file mode 100644
index 0000000000..c03f5cf3d2
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/custom_module.py
@@ -0,0 +1,27 @@
+#!/usr/bin/env python3
+
+import torch
+import torch.cuda.profiler as profiler
+from apex import pyprof
+pyprof.nvtx.init()
+
+class Foo(torch.nn.Module):
+    def __init__(self, size):
+        super(Foo, self).__init__()
+        self.n = torch.nn.Parameter(torch.ones(size))
+        self.m = torch.nn.Parameter(torch.ones(size))
+
+    def forward(self, input):
+        return self.n*input + self.m
+
+#Hook the forward function to pyprof
+pyprof.nvtx.wrap(Foo, 'forward')
+
+foo = Foo(4)
+foo.cuda()
+x = torch.ones(4).cuda()
+
+with torch.autograd.profiler.emit_nvtx():
+	profiler.start()
+	z = foo(x)
+	profiler.stop()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/test.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/test.sh
new file mode 100644
index 0000000000..dc02f388d5
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/custom_func_module/test.sh
@@ -0,0 +1,30 @@
+#!/bin/bash
+
+set -e
+
+SCRIPT=`realpath $0`
+SCRIPTPATH=`dirname $SCRIPT`
+PYPROF="$SCRIPTPATH/../.."
+
+parse="python $PYPROF/parse/parse.py"
+prof="python $PYPROF/prof/prof.py"
+
+for f in *.py
+do
+	base=`basename $f .py`
+	sql=$base.sql
+	dict=$base.dict
+
+	#NVprof
+	echo "nvprof -fo $sql python $f"
+	nvprof -fo $sql python $f
+
+	#Parse
+	echo $parse $sql
+	$parse $sql > $dict
+
+	#Prof
+	echo $prof $dict
+	$prof -w 130 $dict
+	\rm $sql $dict
+done
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/imagenet/imagenet.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/imagenet/imagenet.py
new file mode 100644
index 0000000000..79cd4ecf97
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/imagenet/imagenet.py
@@ -0,0 +1,135 @@
+#!/usr/bin/env python3
+
+"""
+Example to run pyprof with imagenet models.
+"""
+
+import sys
+import torch
+import torch.nn as nn
+import torchvision.models as models
+import torch.cuda.profiler as profiler
+import argparse
+
+from apex import pyprof
+from apex.optimizers import FusedAdam
+
+def parseArgs():
+	parser = argparse.ArgumentParser(prog=sys.argv[0], description="Run popular imagenet models.")
+
+	parser.add_argument("-m",
+		type=str,
+		default="resnet50",
+		choices=["alexnet", "densenet121", "densenet161", "densenet169", "densenet201", "googlenet", "mnasnet0_5", "mnasnet0_75", "mnasnet1_0", "mnasnet1_3", "mobilenet_v2", "resnet18", "resnet34", "resnet50", "resnet101", "resnet152", "resnext50_32x4d", "resnext101_32x8d", "wide_resnet50_2", "wide_resnet101_2", "shufflenet_v2_x0_5", "shufflenet_v2_x1_0", "shufflenet_v2_x1_5", "shufflenet_v2_x2_0", "squeezenet1_0", "squeezenet1_1", "vgg11", "vgg11_bn", "vgg13", "vgg13_bn", "vgg16", "vgg16_bn", "vgg19", "vgg19_bn", "inception_v3"],
+		help="Model.")
+
+	parser.add_argument("-b",
+		type=int,
+		default=32,
+		help="Batch size.")
+
+	parser.add_argument("-o",
+		type=str,
+		default="adam",
+		choices=["adam", "sgd"],
+		help="Optimizer.")
+
+	args = parser.parse_args()
+	return args
+
+d = {
+	"alexnet":				{'H': 224, 'W': 224, 'opts': {}},
+
+	"densenet121":			{'H': 224, 'W': 224, 'opts': {}},
+	"densenet161":			{'H': 224, 'W': 224, 'opts': {}},
+	"densenet169":			{'H': 224, 'W': 224, 'opts': {}},
+	"densenet201":			{'H': 224, 'W': 224, 'opts': {}},
+
+	"googlenet":			{'H': 224, 'W': 224, 'opts': {'aux_logits': False}},
+
+	"mnasnet0_5":			{'H': 224, 'W': 224, 'opts': {}},
+	"mnasnet0_75":			{'H': 224, 'W': 224, 'opts': {}},
+	"mnasnet1_0":			{'H': 224, 'W': 224, 'opts': {}},
+	"mnasnet1_3":			{'H': 224, 'W': 224, 'opts': {}},
+
+	"mobilenet_v2":			{'H': 224, 'W': 224, 'opts': {}},
+
+	"resnet18":				{'H': 224, 'W': 224, 'opts': {}},
+	"resnet34":				{'H': 224, 'W': 224, 'opts': {}},
+	"resnet50":				{'H': 224, 'W': 224, 'opts': {}},
+	"resnet101":			{'H': 224, 'W': 224, 'opts': {}},
+	"resnet152":			{'H': 224, 'W': 224, 'opts': {}},
+
+	"resnext50_32x4d":		{'H': 224, 'W': 224, 'opts': {}},
+	"resnext101_32x8d":		{'H': 224, 'W': 224, 'opts': {}},
+
+	"wide_resnet50_2":		{'H': 224, 'W': 224, 'opts': {}},
+	"wide_resnet101_2":		{'H': 224, 'W': 224, 'opts': {}},
+
+	"shufflenet_v2_x0_5": 	{'H': 224, 'W': 224, 'opts': {}},
+	"shufflenet_v2_x1_0": 	{'H': 224, 'W': 224, 'opts': {}},
+	"shufflenet_v2_x1_5": 	{'H': 224, 'W': 224, 'opts': {}},
+	"shufflenet_v2_x2_0":	{'H': 224, 'W': 224, 'opts': {}},
+
+	"squeezenet1_0":		{'H': 224, 'W': 224, 'opts': {}},
+	"squeezenet1_1":		{'H': 224, 'W': 224, 'opts': {}},
+
+	"vgg11":				{'H': 224, 'W': 224, 'opts': {}},
+	"vgg11_bn":				{'H': 224, 'W': 224, 'opts': {}},
+	"vgg13":				{'H': 224, 'W': 224, 'opts': {}},
+	"vgg13_bn":				{'H': 224, 'W': 224, 'opts': {}},
+	"vgg16":				{'H': 224, 'W': 224, 'opts': {}},
+	"vgg16_bn":				{'H': 224, 'W': 224, 'opts': {}},
+	"vgg19":				{'H': 224, 'W': 224, 'opts': {}},
+	"vgg19_bn":				{'H': 224, 'W': 224, 'opts': {}},
+
+	"inception_v3":			{'H': 299, 'W': 299, 'opts': {'aux_logits': False}},
+	}
+
+def main():
+	args = parseArgs()
+
+	pyprof.nvtx.init()
+#	pyprof.nvtx.wrap(fused_adam_cuda, 'adam')
+
+	N = args.b
+	C = 3
+	H = d[args.m]['H']
+	W = d[args.m]['W']
+	opts = d[args.m]['opts']
+	classes = 1000
+
+	net = getattr(models, args.m)
+	net = net(**opts).cuda().half()
+	net.train()
+
+	x = torch.rand(N, C, H, W).cuda().half()
+	target = torch.empty(N, dtype=torch.long).random_(classes).cuda()
+
+	criterion = nn.CrossEntropyLoss().cuda()
+	if (args.o == "sgd"):
+		optimizer = torch.optim.SGD(net.parameters(), lr = 0.01, momentum=0.9)
+	elif (args.o == "adam"):
+		optimizer = FusedAdam(net.parameters())
+	else:
+		assert False
+
+	#Warm up without profiler
+	for i in range(2):
+		output = net(x)
+		loss = criterion(output, target)
+		optimizer.zero_grad()
+		loss.backward()
+		optimizer.step()
+
+	with torch.autograd.profiler.emit_nvtx():
+		profiler.start()
+		output = net(x)
+		loss = criterion(output, target)
+		optimizer.zero_grad()
+		loss.backward()
+		optimizer.step()
+		profiler.stop()
+
+if __name__ == "__main__":
+	main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/imagenet/test.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/imagenet/test.sh
new file mode 100644
index 0000000000..0c44c05bb8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/imagenet/test.sh
@@ -0,0 +1,36 @@
+#!/bin/bash
+
+set -e
+
+SCRIPT=`realpath $0`
+SCRIPTPATH=`dirname $SCRIPT`
+PYPROF="$SCRIPTPATH/../.."
+
+parse="python -m apex.pyprof.parse"
+prof="python -m apex.pyprof.prof"
+
+for net in "resnet50"
+do
+	for optim in adam sgd
+	do
+		for batch in 32 64
+		do
+			base="torchvision".$net.$optim.$batch
+			sql=$base.sql
+			dict=$base.dict
+
+			#NVprof
+			echo "nvprof -fo $sql --profile-from-start off python imagenet.py -m ${net} -o $optim -b $batch"
+			nvprof -fo $sql --profile-from-start off python imagenet.py -m ${net} -o $optim -b $batch
+
+			#Parse
+			echo $parse $sql
+			$parse $sql > $dict
+
+			#Prof
+			echo $prof $dict
+			$prof -w 130 $dict
+#			\rm $sql $dict
+		done
+	done
+done
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/README.md
new file mode 100644
index 0000000000..101980f007
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/README.md
@@ -0,0 +1,14 @@
+*As of this writing, these examples do not work
+because of changes being proposed in PyTorch.*
+
+There are two ways to use PyTorch JIT
+ - Scripting
+ - Tracing
+
+In addition, we can JIT a
+ - Stand alone function
+ - Class / class method
+
+This directory has an example for each of the 4 cases.
+Intercepting (monkey patching) JITted code has a few extra steps,
+which are explained through comments.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_script_function.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_script_function.py
new file mode 100644
index 0000000000..0b692c8105
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_script_function.py
@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+
+import torch
+import torch.cuda.profiler as profiler
+from apex import pyprof
+
+#The following creates an object "foo" of type ScriptModule
+#The new object has a function called "forward"
+
+@torch.jit.script
+def foo(x, y):
+	return torch.sigmoid(x) + y
+
+#Initialize pyprof after the JIT step
+pyprof.nvtx.init()
+
+#Assign a name to the object "foo"
+foo.__name__ = "foo"
+
+#Hook up the forward function to pyprof
+pyprof.nvtx.wrap(foo, 'forward')
+
+x = torch.zeros(4,4).cuda()
+y = torch.ones(4,4).cuda()
+
+with torch.autograd.profiler.emit_nvtx():
+	profiler.start()
+	z = foo(x, y)
+	profiler.stop()
+	print(z)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_script_method.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_script_method.py
new file mode 100644
index 0000000000..79189c837e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_script_method.py
@@ -0,0 +1,31 @@
+#!/usr/bin/env python3
+
+import torch
+import torch.cuda.profiler as profiler
+from apex import pyprof
+
+class Foo(torch.jit.ScriptModule):
+    def __init__(self, size):
+        super(Foo, self).__init__()
+        self.n = torch.nn.Parameter(torch.ones(size))
+        self.m = torch.nn.Parameter(torch.ones(size))
+
+    @torch.jit.script_method
+    def forward(self, input):
+        return self.n*input + self.m
+
+#Initialize pyprof after the JIT step
+pyprof.nvtx.init()
+
+#Hook up the forward function to pyprof
+pyprof.nvtx.wrap(Foo, 'forward')
+
+foo = Foo(4)
+foo.cuda()
+x = torch.ones(4).cuda()
+
+with torch.autograd.profiler.emit_nvtx():
+	profiler.start()
+	z = foo(x)
+	profiler.stop()
+	print(z)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_trace_function.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_trace_function.py
new file mode 100644
index 0000000000..06a781888f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_trace_function.py
@@ -0,0 +1,30 @@
+#!/usr/bin/env python3
+
+import torch
+import torch.cuda.profiler as profiler
+from apex import pyprof
+
+def foo(x, y):
+	return torch.sigmoid(x) + y
+
+x = torch.zeros(4,4).cuda()
+y = torch.ones(4,4).cuda()
+
+#JIT the function using tracing
+#This returns an object of type ScriptModule with a forward method.
+traced_foo = torch.jit.trace(foo, (x,y))
+
+#Initialize pyprof after the JIT step
+pyprof.nvtx.init()
+
+#Assign a name to the object "traced_foo"
+traced_foo.__dict__['__name__'] = "foo"
+
+#Hook up the forward function to pyprof
+pyprof.nvtx.wrap(traced_foo, 'forward')
+
+with torch.autograd.profiler.emit_nvtx():
+	profiler.start()
+	z = traced_foo(x, y)
+	profiler.stop()
+	print(z)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_trace_method.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_trace_method.py
new file mode 100644
index 0000000000..8a0d3fcecc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/jit_trace_method.py
@@ -0,0 +1,36 @@
+#!/usr/bin/env python3
+
+import torch
+import torch.cuda.profiler as profiler
+from apex import pyprof
+
+class Foo(torch.nn.Module):
+    def __init__(self, size):
+        super(Foo, self).__init__()
+        self.n = torch.nn.Parameter(torch.ones(size))
+        self.m = torch.nn.Parameter(torch.ones(size))
+
+    def forward(self, input):
+        return self.n*input + self.m
+
+foo = Foo(4)
+foo.cuda()
+x = torch.ones(4).cuda()
+
+#JIT the class using tracing
+traced_foo = torch.jit.trace(foo, x)
+
+#Initialize pyprof after the JIT step
+pyprof.nvtx.init()
+
+#Assign a name to the object "traced_foo"
+traced_foo.__dict__['__name__'] = "foo"
+
+#Hook up the forward function to pyprof
+pyprof.nvtx.wrap(traced_foo, 'forward')
+
+with torch.autograd.profiler.emit_nvtx():
+	profiler.start()
+	z = traced_foo(x)
+	profiler.stop()
+	print(z)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/test.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/test.sh
new file mode 100644
index 0000000000..dc02f388d5
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/jit/test.sh
@@ -0,0 +1,30 @@
+#!/bin/bash
+
+set -e
+
+SCRIPT=`realpath $0`
+SCRIPTPATH=`dirname $SCRIPT`
+PYPROF="$SCRIPTPATH/../.."
+
+parse="python $PYPROF/parse/parse.py"
+prof="python $PYPROF/prof/prof.py"
+
+for f in *.py
+do
+	base=`basename $f .py`
+	sql=$base.sql
+	dict=$base.dict
+
+	#NVprof
+	echo "nvprof -fo $sql python $f"
+	nvprof -fo $sql python $f
+
+	#Parse
+	echo $parse $sql
+	$parse $sql > $dict
+
+	#Prof
+	echo $prof $dict
+	$prof -w 130 $dict
+	\rm $sql $dict
+done
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/lenet.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/lenet.py
new file mode 100644
index 0000000000..9e22477eaf
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/lenet.py
@@ -0,0 +1,65 @@
+#!/usr/bin/env python3
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.cuda.profiler as profiler
+import torch.optim as optim
+
+from apex import pyprof
+pyprof.nvtx.init()
+
+class LeNet5(nn.Module):
+	def __init__(self):
+		super(LeNet5, self).__init__()
+		# 1 input image channel, 6 output channels, 5x5 square convolution
+		# kernel
+		self.conv1 = nn.Conv2d(1, 6, 5)
+		self.conv2 = nn.Conv2d(6, 16, 5)
+		# an affine operation: y = Wx + b
+		self.fc1 = nn.Linear(16 * 5 * 5, 120)
+		self.fc2 = nn.Linear(120, 84)
+		self.fc3 = nn.Linear(84, 10)
+
+	def forward(self, x):
+		# Max pooling over a (2, 2) window
+		x = F.max_pool2d(F.relu(self.conv1(x)), (2, 2))
+		# If the size is a square you can only specify a single number
+		x = F.max_pool2d(F.relu(self.conv2(x)), 2)
+		x = x.view(-1, self.num_flat_features(x))
+		x = F.relu(self.fc1(x))
+		x = F.relu(self.fc2(x))
+		x = self.fc3(x)
+		return x
+
+	def num_flat_features(self, x):
+		size = x.size()[1:]  # all dimensions except the batch dimension
+		num_features = 1
+		for s in size:
+			num_features *= s
+		return num_features
+
+with torch.autograd.profiler.emit_nvtx():
+
+	net = LeNet5().cuda()
+
+	input = torch.randn(1, 1, 32, 32).cuda()
+	out = net(input)
+
+	target = torch.randn(10)			# a dummy target, for example
+	target = target.view(1, -1).cuda()	# make it the same shape as output
+	criterion = nn.MSELoss()
+
+	# create your optimizer
+	optimizer = optim.SGD(net.parameters(), lr=0.01)
+
+	# in your training loop:
+	optimizer.zero_grad()	# zero the gradient buffers
+
+	profiler.start()
+	output = net(input)
+	loss = criterion(output, target)
+	loss.backward()
+	optimizer.step()	# Does the update
+	profiler.stop()
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/operators.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/operators.py
new file mode 100644
index 0000000000..a033c2ef47
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/operators.py
@@ -0,0 +1,145 @@
+#!/usr/bin/env python3
+
+"""
+This file checks all Python operators.
+"""
+
+import sys
+import torch
+import torch.cuda.profiler as profiler
+import operator
+import inspect
+
+#Import and initialize pyprof
+from apex import pyprof
+pyprof.nvtx.init()
+
+X = 1024
+Y = 1024
+
+fa = torch.rand(X, Y).cuda()
+fb = torch.rand(X, Y).cuda()
+fc = torch.rand(X, Y).cuda()
+
+ia = torch.randint(0, 100, (X, Y)).cuda()
+ib = torch.randint(0, 100, (X, Y)).cuda()
+
+sa = torch.ones(1,1).cuda()
+sb = torch.ones(1,1).cuda()
+
+ba = fa.byte()
+
+unaryOps = ["abs", "__abs__", "neg", "__neg__",]
+invertOps = ["inv", "invert", "__inv__", "__invert__",]	#imlemented only for byte tensors
+#pos, __pos__ is not implemented for tensors
+
+binaryOps = []
+binaryOps += [ "lt", "__lt__", "le", "__le__", "eq", "__eq__", "ne", "__ne__", "ge", "__ge__", "gt", "__gt__" ]
+binaryOps += [ "add", "__add__", "sub", "__sub__", "mul", "__mul__", "floordiv", "__floordiv__", "truediv", "__truediv__", "pow", "__pow__", "mod", "__mod__"]
+binaryOps += [ "and_", "__and__", "or_", "__or__", "xor", "__xor__", "lshift", "__lshift__", "rshift", "__rshift__"]
+
+inplaceOps = []
+inplaceOps += ["iadd", "__iadd__", "isub", "__isub__", "imul", "__imul__", "ifloordiv", "__ifloordiv__", "itruediv", "__itruediv__", "imod", "__imod__",]
+#ipow, __ipow__ is not implemented in pytorch
+inplaceOps += [ "iand", "__iand__", "ior", "__ior__", "ixor", "__ixor__", "ilshift", "__ilshift__", "irshift", "__irshift__",]
+
+matmulOps = [ "matmul", "__matmul__" ]
+inplacematmulOps = [ "imatmul", "__imatmul__" ]
+
+reverseIntBinaryOps = ["__radd__", "__rsub__", "__rmul__", "__rfloordiv__", "__rpow__",]
+reverseFloatBinaryOps = ["__radd__", "__rsub__", "__rmul__", "__rdiv__", "__rtruediv__", "__rfloordiv__", "__rpow__",]
+
+'''
+TODO
+.concat(a, b)
+.__concat__(a, b)
+.contains(a, b)
+.__contains__(a, b)
+.countOf(a, b)
+.delitem(a, b)
+.__delitem__(a, b)
+.getitem(a, b)
+.__getitem__(a, b)
+.indexOf(a, b)
+.setitem(a, b, c)
+.__setitem__(a, b, c)
+.length_hint(obj, default=0)
+.iconcat(a, b)
+.__iconcat__(a, b)
+.index(a)
+.__index__(a)
+'''
+
+#Context manager
+with torch.autograd.profiler.emit_nvtx():
+
+	#Start profiler
+	profiler.start()
+
+	for op in unaryOps:
+		assert hasattr(operator, op)
+		f = getattr(operator, op)
+		assert inspect.isbuiltin(f)
+		c = f(ia)
+
+	for op in invertOps:
+		assert hasattr(operator, op)
+		f = getattr(operator, op)
+		assert inspect.isbuiltin(f)
+		c = f(ba)
+
+	for op in binaryOps:
+		assert hasattr(operator, op)
+		f = getattr(operator, op)
+		assert inspect.isbuiltin(f)
+		c = f(ia, ib)
+		c = f(ia, 2)
+
+	for op in inplaceOps:
+		assert hasattr(operator, op)
+		f = getattr(operator, op)
+		assert inspect.isbuiltin(f)
+		ia = f(ia, ib)
+		ia = f(ia, 2)
+
+	for op in matmulOps:
+		assert hasattr(operator, op)
+		f = getattr(operator, op)
+		assert inspect.isbuiltin(f)
+		c = f(fa, fb)
+
+	for op in inplacematmulOps:
+		assert hasattr(operator, op)
+		f = getattr(operator, op)
+		assert inspect.isbuiltin(f)
+		fa = f(fa, fb)
+
+	for op in reverseIntBinaryOps:
+		assert hasattr(torch.Tensor, op)
+		f = getattr(torch.Tensor, op)
+		ia = f(ia, ib)
+
+	for op in reverseFloatBinaryOps:
+		assert hasattr(torch.Tensor, op)
+		f = getattr(torch.Tensor, op)
+		fa = f(fa, fb)
+
+	'''
+	#c = fa[3]
+	#c = fa[3][3]
+	#c = torch.min(fa, 3)
+	c = torch.sum(fa)
+	c = torch.max(fa)
+	c = -fa
+	#fc[2][2] = fa[2][2]
+
+	c = a_scalar and b_scalar
+	c = a_scalar or b_scalar
+	c = not a_scalar
+
+	c = a is b
+	c = a is not b
+	'''
+
+	#Stop profiler
+	profiler.stop()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/simple.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/simple.py
new file mode 100644
index 0000000000..dbe9c615fa
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/simple.py
@@ -0,0 +1,38 @@
+#!/usr/bin/env python3
+
+"""
+This simple file provides an example of how to
+ - import the pyprof library and initialize it
+ - use the emit_nvtx context manager
+ - start and stop the profiler
+
+Only kernels within profiler.start and profiler.stop calls are profiled.
+To profile
+$ nvprof -f -o simple.sql --profile-from-start off ./simple.py
+"""
+
+import sys
+import torch
+import torch.cuda.profiler as profiler
+
+#Import and initialize pyprof
+from apex import pyprof
+pyprof.nvtx.init()
+
+a = torch.randn(5, 5).cuda()
+b = torch.randn(5, 5).cuda()
+
+#Context manager
+with torch.autograd.profiler.emit_nvtx():
+
+	#Start profiler
+	profiler.start()
+
+	c = a + b
+	c = torch.mul(a,b)
+	c = torch.matmul(a,b)
+	c = torch.argmax(a, dim=1)
+	c = torch.nn.functional.pad(a, (1,1))
+
+	#Stop profiler
+	profiler.stop()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/README.md
new file mode 100644
index 0000000000..30523615a3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/README.md
@@ -0,0 +1,21 @@
+Nvidia NVTX range markers (https://docs.nvidia.com/gameworks/content/gameworkslibrary/nvtx/nvidia_tools_extension_library_nvtx.htm) 
+are a useful tool to capture and observe events and code ranges etc. 
+Using PyTorch APIs e.g, `torch.cuda.nvtx.range_push("xxx")` and `torch.cuda.nvtx.range_pop()` users can easily add their own NVTX range markers. These markers can then be observed in the Nvidia Visual Profiler (NVVP).
+
+While inserting NVTX markers (strings), if the users follow a specific string pattern `"layer:your_string_here"` e.g. `"layer:conv1"` or `"layer:encoder_layer_3_self_attention`, then `pyprof` will display the strings `conv1` and `encoder_layer_3_self_attention` next to the associated kernels in the output of `prof.py` when used with the `-c layer` option.
+
+NVTX range markers can be nested and if users follow the above string pattern, the output of `prof.py` will show all the markers associated with a kernel.
+
+The file `resnet.py` (a simplified version of the torchvision model) shows an example of how users can add (nested) NVTX markers with information which can greatly aid in understanding and analysis of networks.
+
+Note that the pattern `"layer:your_string_here"` was chosen to aid information extraction by `pyprof`. The tool will work seamlessly even if there are other markers or no markers at all.
+
+### To run
+
+```sh
+nvprof -fo resnet.sql --profile-from-start off python resnet.py
+parse.py resnet.sql > resnet.dict
+prof.py --csv -c idx,layer,dir,mod,op,kernel,params,sil resnet.dict
+```
+
+The file `resnet.sql` can also be opened with NVVP as usual.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/resnet.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/resnet.py
new file mode 100644
index 0000000000..ec351aa974
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/resnet.py
@@ -0,0 +1,215 @@
+#!/usr/bin/env python3
+
+"""
+An example showing use of nested NVTX markers.
+"""
+
+import torch
+import torch.nn as nn
+
+import torch.cuda.profiler as profiler
+import torch.cuda.nvtx as nvtx
+from apex import pyprof
+pyprof.nvtx.init()
+
+def conv3x3(in_planes, out_planes, stride=1, groups=1, dilation=1):
+	"""3x3 convolution with padding"""
+	return nn.Conv2d(in_planes, out_planes, kernel_size=3, stride=stride,
+					 padding=dilation, groups=groups, bias=False, dilation=dilation)
+
+def conv1x1(in_planes, out_planes, stride=1):
+	"""1x1 convolution"""
+	return nn.Conv2d(in_planes, out_planes, kernel_size=1, stride=stride, bias=False)
+
+class Bottleneck(nn.Module):
+	expansion = 4
+	count = 1
+
+	def __init__(self, inplanes, planes, stride=1, downsample=None, groups=1,
+				 base_width=64, dilation=1, norm_layer=None):
+		super(Bottleneck, self).__init__()
+		if norm_layer is None:
+			norm_layer = nn.BatchNorm2d
+		width = int(planes * (base_width / 64.)) * groups
+		# Both self.conv2 and self.downsample layers downsample the input when stride != 1
+		self.conv1 = conv1x1(inplanes, width)
+		self.bn1 = norm_layer(width)
+		self.conv2 = conv3x3(width, width, stride, groups, dilation)
+		self.bn2 = norm_layer(width)
+		self.conv3 = conv1x1(width, planes * self.expansion)
+		self.bn3 = norm_layer(planes * self.expansion)
+		self.relu = nn.ReLU(inplace=True)
+		self.downsample = downsample
+		self.stride = stride
+
+		self.id = Bottleneck.count
+		Bottleneck.count += 1
+
+	def forward(self, x):
+		identity = x
+
+		nvtx.range_push("layer:Bottleneck_{}".format(self.id))
+
+		nvtx.range_push("layer:Conv1")
+		out = self.conv1(x)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:BN1")
+		out = self.bn1(out)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:ReLU")
+		out = self.relu(out)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:Conv2")
+		out = self.conv2(out)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:BN2")
+		out = self.bn2(out)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:ReLU")
+		out = self.relu(out)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:Conv3")
+		out = self.conv3(out)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:BN3")
+		out = self.bn3(out)
+		nvtx.range_pop()
+
+		if self.downsample is not None:
+			nvtx.range_push("layer:Downsample")
+			identity = self.downsample(x)
+			nvtx.range_pop()
+
+		nvtx.range_push("layer:Residual")
+		out += identity
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:ReLU")
+		out = self.relu(out)
+		nvtx.range_pop()
+
+		nvtx.range_pop()
+
+		return out
+
+class ResNet(nn.Module):
+
+	def __init__(self, block, layers, num_classes=1000,
+				 groups=1, width_per_group=64, norm_layer=None):
+		super(ResNet, self).__init__()
+		if norm_layer is None:
+			norm_layer = nn.BatchNorm2d
+		self._norm_layer = norm_layer
+
+		self.inplanes = 64
+		self.dilation = 1
+
+		self.groups = groups
+		self.base_width = width_per_group
+		self.conv1 = nn.Conv2d(3, self.inplanes, kernel_size=7, stride=2, padding=3, bias=False)
+		self.bn1 = norm_layer(self.inplanes)
+		self.relu = nn.ReLU(inplace=True)
+		self.maxpool = nn.MaxPool2d(kernel_size=3, stride=2, padding=1)
+		self.layer1 = self._make_layer(block, 64, layers[0])
+		self.layer2 = self._make_layer(block, 128, layers[1], stride=2)
+		self.layer3 = self._make_layer(block, 256, layers[2], stride=2)
+		self.layer4 = self._make_layer(block, 512, layers[3], stride=2)
+		self.avgpool = nn.AdaptiveAvgPool2d((1, 1))
+		self.fc = nn.Linear(512 * block.expansion, num_classes)
+
+		for m in self.modules():
+			if isinstance(m, nn.Conv2d):
+				nn.init.kaiming_normal_(m.weight, mode='fan_out', nonlinearity='relu')
+			elif isinstance(m, (nn.BatchNorm2d, nn.GroupNorm)):
+				nn.init.constant_(m.weight, 1)
+				nn.init.constant_(m.bias, 0)
+
+	def _make_layer(self, block, planes, blocks, stride=1, dilate=False):
+		norm_layer = self._norm_layer
+		downsample = None
+		previous_dilation = self.dilation
+		if dilate:
+			self.dilation *= stride
+			stride = 1
+		if stride != 1 or self.inplanes != planes * block.expansion:
+			downsample = nn.Sequential(
+				conv1x1(self.inplanes, planes * block.expansion, stride),
+				norm_layer(planes * block.expansion),
+			)
+
+		layers = []
+		layers.append(block(self.inplanes, planes, stride, downsample, self.groups,
+							self.base_width, previous_dilation, norm_layer))
+		self.inplanes = planes * block.expansion
+		for _ in range(1, blocks):
+			layers.append(block(self.inplanes, planes, groups=self.groups,
+								base_width=self.base_width, dilation=self.dilation,
+								norm_layer=norm_layer))
+
+		return nn.Sequential(*layers)
+
+	def forward(self, x):
+
+		nvtx.range_push("layer:conv1_x")
+		x = self.conv1(x)
+		x = self.bn1(x)
+		x = self.relu(x)
+		x = self.maxpool(x)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:conv2_x")
+		x = self.layer1(x)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:conv3_x")
+		x = self.layer2(x)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:conv4_x")
+		x = self.layer3(x)
+		nvtx.range_pop()
+
+		nvtx.range_push("layer:conv5_x")
+		x = self.layer4(x)
+		nvtx.range_pop()
+
+		x = self.avgpool(x)
+		x = torch.flatten(x, 1)
+
+		nvtx.range_push("layer:FC")
+		x = self.fc(x)
+		nvtx.range_pop()
+
+		return x
+
+
+def resnet50():
+	return ResNet(Bottleneck, [3, 4, 6, 3])
+
+#Create model
+net = resnet50().cuda().half()
+net.train()
+
+#Create optimizer
+criterion = nn.CrossEntropyLoss().cuda()
+optimizer = torch.optim.SGD(net.parameters(), lr = 0.01, momentum=0.9)
+
+#Create synthetic input and label
+x = torch.rand(32, 3, 224, 224).cuda().half()
+target = torch.empty(32, dtype=torch.long).random_(1000).cuda()
+
+with torch.autograd.profiler.emit_nvtx():
+	profiler.start()
+	output = net(x)
+	loss = criterion(output, target)
+	optimizer.zero_grad()
+	loss.backward()
+	optimizer.step()
+	profiler.stop()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/test.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/test.sh
new file mode 100644
index 0000000000..89b8eb0c84
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/examples/user_annotation/test.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+
+set -e
+
+SCRIPT=`realpath $0`
+SCRIPTPATH=`dirname $SCRIPT`
+PYPROF="$SCRIPTPATH/../.."
+
+parse="python $PYPROF/parse/parse.py"
+prof="python $PYPROF/prof/prof.py"
+
+for f in *.py
+do
+	base=`basename $f .py`
+	sql=$base.sql
+	dict=$base.dict
+
+	#NVprof
+	echo "nvprof -fo --profile-from-start off $sql python $f"
+	nvprof -fo $sql --profile-from-start off python $f
+
+	#Parse
+	echo $parse $sql
+	$parse $sql > $dict
+
+	#Prof
+	echo $prof $dict
+	#$prof -w 130 $dict
+	$prof --csv -c idx,layer,dir,mod,op,kernel,params,sil $dict
+	\rm $sql $dict
+done
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/nvtx/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/nvtx/__init__.py
new file mode 100644
index 0000000000..774b4c24fa
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/nvtx/__init__.py
@@ -0,0 +1,2 @@
+from .nvmarker import init
+from .nvmarker import add_wrapper as wrap
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/nvtx/nvmarker.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/nvtx/nvmarker.py
new file mode 100644
index 0000000000..c4d99e3bb9
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/nvtx/nvmarker.py
@@ -0,0 +1,222 @@
+"""
+This file intercepts (monkey patches) the following functions and adds NVTX markers.
+	torch.*
+	torch.Tensor.*
+	torch.nn.functional.*
+	torch.nn.*.forward
+
+The NVTX markers (one or more) contain the following information
+	call trace (a list of file_name:line_number)
+	extra_repr() from torch.nn modules
+	module/class name
+	function name
+	inputs (args and kwargs)
+		scalar: name, type and value
+		tensor: name, shape and datatype
+		numpy: name, shape and datatype
+		list/tuple: a sequence of scalars or tensors or numpy arrays
+"""
+
+import torch
+import torch.cuda.nvtx as nvtx
+import numpy
+import inspect as ins
+import traceback
+import math
+
+def isfunc(mod, f):
+	assert hasattr(mod, f)
+	attr = getattr(mod, f)
+
+	#Ignore functions like _add
+	if (len(f) >= 2):
+		if f[0] == "_" and f[1] != "_":
+			return False
+
+	#Ignore functions from this list
+	ignore = ['__all__', '__array__', '__array_priority__', '__array_wrap__', '__bool__', '__builtins__', '__cached__', '__class__', '__deepcopy__', '__delattr__', '__delitem__', '__dict__', '__dir__', '__doc__', '__file__', '__format__', '__getattribute__', '__getitem__', '__hash__', '__index__', '__init__', '__init_subclass__', '__iter__', '__len__', '__loader__', '__module__', '__name__', '__new__', '__nonzero__', '__package__', '__path__', '__reduce__', '__reduce_ex__', '__repr__', '__reversed__', '__setattr__', '__setitem__', '__setstate__', '__sizeof__', '__spec__', '__str__', '__subclasshook__', '__version__', '__weakref__']
+
+	#Add functions to this list if they cause recursion
+	ignore += ['size', 'tolist', 'dim', 'is_storage', 'item']
+	if f in ignore:
+		return False
+
+	return ins.ismethod(attr) or ins.isfunction(attr) or ins.ismethoddescriptor(attr) or ins.isbuiltin(attr)
+
+def traceMarker(stack):
+	d = {}
+	cadena = []
+	for i in range(len(stack)-1):
+		fi = stack[i]
+		t = "{}:{}".format(fi.filename, fi.lineno)
+		cadena.append(t)
+	d['traceMarker'] = cadena
+	return str(d)
+
+def modMarker(mod, fn_name, args):
+	"""
+	Returns the stringified extra_repr() of a module.
+	"""
+	assert(fn_name == 'forward')
+	assert(len(args) > 0)
+	d = {}
+	d['mod'] = mod.__name__
+	d['strRepr'] = args[0].extra_repr()
+	return str(d)
+
+def add_wrapper(mod, fn_name):
+	assert isfunc(mod, fn_name)
+
+	# Get a pointer to the original function
+	func = getattr(mod, fn_name)
+
+	# Check if the mod has a string representation
+	# and is not a Script or Traced module (used by JIT)
+	s = hasattr(mod, "extra_repr") and (type(mod) is not torch.jit.ScriptModule) and (type(mod) is not torch.jit.TopLevelTracedModule)
+
+	def wrapper_func(*args, **kwargs):
+
+		# Extract the stacktrace
+		stack = traceback.extract_stack()
+
+		# Push trace marker
+		nvtx.range_push(traceMarker(stack))
+
+		# Push module marker
+		if s:
+			m = modMarker(mod, fn_name, args)
+			nvtx.range_push(m)
+
+		# Create and push argument marker
+		cadena = argMarker(mod, fn_name, args, kwargs)
+		nvtx.range_push(cadena)
+
+		# Call the original function
+		result = func(*args, **kwargs)
+
+		# Pop argumet marker
+		nvtx.range_pop()
+
+		# Pop module marker
+		if s:
+			nvtx.range_pop()
+
+		# Pop trace marker
+		nvtx.range_pop()
+
+		return result
+	setattr(mod, fn_name, wrapper_func)
+
+def argMarker(mod, op, args, kwargs):
+	#For this function args is a tuple and kwargs is a dict
+
+	def tensor(arg, name=""):
+		a = {}
+		a['name'] = name
+		a['type'] = "tensor"
+		a['shape'] = tuple(arg.size())
+		a['dtype'] = str(arg.dtype).split(".")[-1]
+		cadena['args'].append(a)
+
+	def ndarray(arg, name=""):
+		a = {}
+		a['name'] = name
+		a['type'] = "ndarray"
+		a['shape'] = arg.shape
+		a['dtype'] = str(arg.dtype).split(".")[-1]
+		cadena['args'].append(a)
+
+	def seq(arg, name=""):
+		assert issequence(arg)
+		a = {}
+		a['name'] = name
+		if isinstance(arg, list):
+			a['type'] = "list"
+			a['value'] = arg
+		else:
+			a['type'] = "tuple"
+			# The arg could be torch.Size, which is a subclass of tuple
+			# Therefore, explicitly convert to tuple
+			a['value'] = tuple(arg)
+		
+		cadena['args'].append(a)
+
+	def scalar(arg, name=""):
+		a = {}
+		a['name'] = name
+		a['type'] = type(arg).__name__
+		#handle the case when the argument is +/- inf or nan
+		if arg == float('inf'):
+			a['value'] = "inf"
+		elif arg == float('-inf'):
+			a['value'] = "-inf"
+		elif isinstance(arg, float) and math.isnan(arg):
+			a['value'] = "nan"
+		else:
+			a['value'] = arg
+		cadena['args'].append(a)
+
+	def isscalar(arg):
+		return (type(arg) is int) or (type(arg) is float) or (type(arg) is bool) or (arg is None) or (type(arg) is str)
+
+	def issequence(arg):
+		return isinstance(arg, list) or isinstance(arg, tuple)
+
+	def foo(args, name):
+		#args should be an iterable sequence e.g. list or tuple
+		for arg in args:
+			if isinstance(arg, torch.Tensor):
+				if arg.dim() == 0:
+					scalar(arg.item(), name)
+				else:
+					tensor(arg, name)
+			elif isinstance(arg, numpy.ndarray):
+				ndarray(arg, name)
+			elif (isscalar(arg)):
+				scalar(arg, name)
+			elif issequence(arg):
+				if (len(arg) == 0) or isscalar(arg[0]):	#An empty sequence or a sequence of scalars
+					seq(arg, name)
+				else:	# A sequence of tensors or numpy arrays
+					foo(arg, name)
+			'''
+			else:
+				print("The following arg is none of Tensor, numpy array, scalar but a %s" % (str(type(arg))))
+				print("Mod: %s" % str(mod.__name__))
+				print("Op: %s" % str(op))
+				print(dir(arg))
+			'''
+
+	cadena = {}
+	cadena['mod'] = mod.__name__
+	cadena['op'] = op
+	cadena['args'] = []
+
+	foo(args, "")
+	for k,v in kwargs.items():
+		foo((v,), k)
+
+	return str(cadena)
+
+def patchClass(cls):
+	for f in dir(cls):
+		if isfunc(cls, f):
+			add_wrapper(cls, f)
+
+def init():
+	string = "\n\nPyprof has been moved to its own dedicated repository and will " + \
+			"soon be removed from Apex.  Please visit\n" + \
+			"https://github.com/NVIDIA/PyProf\n" + \
+			"for the latest version.\n\n"
+	# print regardless of warning state
+	print(string)
+
+	print("Initializing NVTX monkey patches")
+	for cls in [torch, torch.Tensor, torch.nn.functional,]:
+		patchClass(cls)
+
+	for cls in [torch.nn.RNN, torch.nn.RNNCell, torch.nn.LSTM, torch.nn.LSTMCell, torch.nn.GRU, torch.nn.GRUCell]:
+		if isfunc(cls, 'forward'):
+			add_wrapper(cls, 'forward')
+
+	print("Done with NVTX monkey patching")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/__main__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/__main__.py
new file mode 100644
index 0000000000..fa12919923
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/__main__.py
@@ -0,0 +1,10 @@
+import warnings
+
+try:
+    from .parse import main
+except ImportError as e:
+    warnings.warn("Did you make sure to install PyProf dependencies by using the --pyprof flag during Apex installation?)")
+    raise e
+
+if __name__ == '__main__':
+    main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/db.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/db.py
new file mode 100644
index 0000000000..55c5dc5714
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/db.py
@@ -0,0 +1,61 @@
+import sys, sqlite3
+
+class DB(object):
+	"""
+	This class provides functions for DB operations
+	with exception handling.
+	"""
+
+	def __init__(self, dbFile):
+		try:
+			conn = sqlite3.connect(dbFile)
+			conn.row_factory = sqlite3.Row
+			c = conn.cursor()
+		except:
+			print("Error opening {}".format(dbFile))
+			sys.exit(1)
+
+		self.conn = conn
+		self.c = c
+
+	def select(self, cmd):
+		try:
+			self.c.execute(cmd)
+			#rows = self.c.fetchall()
+			rows = [dict(row) for row in self.c.fetchall()]
+		except sqlite3.Error as e:
+			print(e)
+			sys.exit(1)
+		except:
+			print("Uncaught error in SQLite access while executing {}".format(cmd))
+			sys.exit(1)
+
+		#print(rows)
+		return rows
+
+	def insert(self, cmd, data):
+		try:
+			self.c.execute(cmd, data)
+		except sqlite3.Error as e:
+			print(e)
+			sys.exit(1)
+		except:
+			print("Uncaught error in SQLite access while executing {}".format(cmd))
+			sys.exit(1)
+
+	def execute(self, cmd):
+		try:
+			self.c.execute(cmd)
+		except sqlite3.Error as e:
+			print(e)
+			sys.exit(1)
+		except:
+			print("Uncaught error in SQLite access while executing {}".format(cmd))
+			sys.exit(1)
+
+	def commit(self):
+		self.conn.commit()
+
+	def close(self):
+		self.c.close()
+		self.conn.close()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/kernel.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/kernel.py
new file mode 100644
index 0000000000..5a66e290d2
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/kernel.py
@@ -0,0 +1,210 @@
+import cxxfilt, struct, binascii
+
+#Helper functions
+
+def demangle(name):
+	"""
+	Demangle a C++ string
+	"""
+	return cxxfilt.demangle(name)
+
+def encode_object_id(pid, tid):
+	"""
+	Given process id (pid) and thread id (tid), return the object id.
+	object id = pid (little endian 4 bytes) + tid (little endian 8 bytes)
+	"""
+	objId = struct.pack('<i', pid) + struct.pack('<q',tid)
+	objId = binascii.hexlify(objId).decode('ascii').upper()
+	return objId
+
+def getShortName(name):
+	"""
+	Returns a shorter kernel name
+	"""
+	sname = name.split("<")[0] \
+				.replace("void ", "") \
+				.replace("at::","") \
+				.replace("cuda::", "") \
+				.replace("native::","") \
+				.replace("(anonymous namespace)::", "")
+	sname = sname.split("(")[0]
+	return sname
+
+class Kernel(object):
+	"""
+	This class stores information about a kernel.
+	"""
+
+	kernels = []
+	profStart = 0
+
+	def __init__(self):
+		self.kNameId = None
+		self.kShortName = None
+		self.kLongName = None
+		self.kStartTime = None	#GPU start time
+		self.kEndTime = None	#GPU end time
+		self.kDuration = None
+		self.device = None
+		self.stream = None
+		self.grid = ()
+		self.block = ()
+		self.corrId = None
+		self.rStartTime = None	#CPU start time
+		self.rEndTime = None	#CPU end time
+		self.rDuration = None
+		self.tid = None
+		self.pid = None
+		self.objId = None
+		self.timeOffset = None
+
+		self.layerMarkers = []
+		self.traceMarkers = []
+		self.reprMarkers = []
+		self.pyprofMarkers = []
+		self.seqMarkers = []
+		self.otherMarkers = []
+		self.altMarkers = []
+		self.seqId = []
+		self.altSeqId = []
+		self.layer = []
+
+		self.subSeqId = None
+		self.dir = None
+		self.mod = []
+		self.op = []
+
+	def setKernelInfo(self, info):
+		self.kNameId = info['name']
+		self.corrId = int(info['correlationId'])
+		start = int(info['start'])
+		end = int(info['end'])
+		assert end > start, "This assertion can fail for very large profiles. It usually fails when start = end = 0."
+		self.kStartTime = start
+		self.kEndTime = end
+		self.kDuration = end - start
+		assert (start > Kernel.profStart)
+		self.device = int(info['deviceId'])
+		self.stream = int(info['streamId'])
+		self.grid = (info['gridX'], info['gridY'], info['gridZ'])
+		self.block = (info['blockX'], info['blockY'], info['blockZ'])
+		self.timeOffset = Kernel.profStart
+
+	def setKernelName(self, name):
+		cadena = demangle(name)
+		self.kLongName = cadena
+		self.kShortName = getShortName(cadena)
+
+	def setRunTimeInfo(self, info):
+		start, end, pid, tid = info
+		self.rStartTime = start
+		self.rEndTime = end
+		self.rDuration = end - start
+		self.pid = pid
+		self.tid = tid
+		self.objId = encode_object_id(pid, tid)
+
+	def setMarkerInfo(self, info):
+		self.layerMarkers, self.traceMarkers, self.reprMarkers, self.pyprofMarkers, self.seqMarkers, self.otherMarkers, self.altMarkers, self.seqId, self.altSeqId, self.layer = info
+		self.subSeqId = 0
+
+	def setDirection(self):
+		"""
+		Set direction (fprop, bprop) based on PyTorch sequence markers.
+		It is a heuristic and not a foolproof method.
+		"""
+		if	any("Backward, seq = " in x for x in self.seqMarkers) or \
+			any("backward, seq = " in x for x in self.seqMarkers) or \
+			any("Backward0, seq = " in x for x in self.seqMarkers):
+			self.dir = "bprop"
+		else:
+			self.dir = "fprop"
+
+	def setOp(self):
+		"""
+		Detect and set the class/module (mod) and operation (op)
+		of the kernel e.g. torch.nn.functional / linear, torch / sigmoid.
+		The lookup sequence we use is
+			NVTX markers inserted by pyprof
+			NVTX markers inserted by PyTorch in bprop
+			NVTX markers inserted by PyTorch in fprop
+		It is a heuristic and not a foolproof method.
+		"""
+
+		def sanitize(name):
+			name = name.replace("torch","") \
+						.replace("autograd","") \
+						.replace("_backward","") \
+						.replace("::","") \
+						.replace("jit","") \
+						.replace("(anonymous namespace)","")
+			head, sep, tail = name.partition("Backward")
+			return head
+
+		#Check pyprof markers
+		for m in self.pyprofMarkers:
+			assert ("mod" in m) and ("op" in m) and ("args" in m)
+			t = eval(m)
+			self.op.append(t['op'])
+			self.mod.append(t['mod'])
+
+		if len(self.op):
+			return
+
+		#Check bprop kernel markers
+		for m in self.seqMarkers:
+			if ("backward, seq = " in m) or ("Backward, seq = " in m):
+				op = m.split(",")[0]
+				op = sanitize(op)
+				self.op.append(op)
+				self.mod.append('na')
+
+		if len(self.op):
+			return
+
+		#Check markers with "seq = "
+		for m in self.seqMarkers:
+			if ", seq = " in m:
+				op = m.split(",")[0]
+				self.op.append(op)
+				self.mod.append('na')
+
+		if len(self.op):
+			return
+
+		#If nothing else
+		if len(self.otherMarkers):
+			self.op.append(self.otherMarkers[0])
+		self.mod.append('na')
+
+	def print(self):
+		"""
+		Print kernel information. This is used by prof.py.
+		"""
+
+		a = lambda: None
+		a.kShortName = self.kShortName
+		a.kDuration = self.kDuration
+		#a.layerMarkers = self.layerMarkers
+		a.layer = self.layer
+		a.trace = self.traceMarkers
+		a.reprMarkers = self.reprMarkers
+		a.marker = self.pyprofMarkers
+		a.seqMarker = self.seqMarkers
+
+		a.seqId = self.seqId
+		a.subSeqId = self.subSeqId
+		a.altSeqId = self.altSeqId
+
+		a.dir = self.dir
+		a.mod = self.mod
+		a.op = self.op
+
+		a.tid = self.tid
+		a.device = self.device
+		a.stream = self.stream
+		a.grid = self.grid
+		a.block = self.block
+		a.kLongName = self.kLongName
+
+		print(a.__dict__)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/nvvp.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/nvvp.py
new file mode 100644
index 0000000000..7167a892d6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/nvvp.py
@@ -0,0 +1,282 @@
+import sys
+
+class NVVP(object):
+	"""
+	This class gets kernel information from the SQL (nvvp) database.
+	"""
+
+	driverT = "CUPTI_ACTIVITY_KIND_DRIVER"
+	runtimeT = "CUPTI_ACTIVITY_KIND_RUNTIME"
+	kernelT = "CUPTI_ACTIVITY_KIND_CONCURRENT_KERNEL"
+	markerT = "CUPTI_ACTIVITY_KIND_MARKER"
+	stringT = "StringTable"
+
+	def __init__(self, db):
+		self.db = db
+		self.markerId = 0
+
+	def getProfileStart(self):
+		"""
+		Get the profile start time
+		"""
+		profStart = sys.maxsize
+		for table in [self.driverT, self.runtimeT, self.kernelT, self.markerT]:
+			colname = "timestamp" if table is self.markerT else "start"
+			cmd = "select {} from {} ORDER BY {} ASC LIMIT 1".format(colname, table, colname)
+			result = self.db.select(cmd)
+			assert(len(result) <= 1)
+			if (len(result) == 1):
+				assert(colname in result[0])
+				t = result[0][colname]
+				if (t < profStart):
+					profStart = t
+		assert(profStart < sys.maxsize)
+		return profStart
+
+	def getString(self, id_):
+		"""
+		Get the string associated with an id.
+		"""
+		cmd = "select value from {} where _id_ = {}".format(self.stringT, id_)
+		result = self.db.select(cmd)
+		assert (len(result) == 1)
+		return result[0]['value']
+
+	def createMarkerTable(self):
+		"""
+		Create a temporary table and index it to speed up repeated SQL quesries.
+		The table is an INNER JOIN of CUPTI_ACTIVITY_KIND_MARKER with itself.
+		"""
+		cmd = 'CREATE TEMPORARY TABLE marker AS SELECT \
+					a._id_ as id, \
+					a.timestamp AS startTime, \
+					b.timestamp AS endTime, \
+					HEX(a.objectId) AS objectId, \
+					a.name AS name \
+					FROM {} AS a INNER JOIN {} AS b ON \
+					a.id = b.id and \
+					a.flags = 2 and b.flags = 4'.format(self.markerT, self.markerT)
+		self.db.execute(cmd)
+
+		self.db.execute('CREATE INDEX start_index ON marker (startTime)')
+		self.db.execute('CREATE INDEX end_index ON marker (endTime)')
+		self.db.execute('CREATE INDEX id_index ON marker (id)')
+
+	def getCPUInfo(self, corrId):
+		"""
+		Given the correlation id, get CPU start, end, thread id, process id.
+		The information can be in the runtime table or the driver table.
+		"""
+
+		#First look in the runtime table
+		cmd = "select start,end,processId,threadId from {} where correlationId={}".format(self.runtimeT, corrId);
+		result = self.db.select(cmd)
+		assert (len(result) <= 1)
+
+		if (len(result) == 0):
+			#Look in the driver table
+			cmd = "select start,end,processId,threadId from {} where correlationId={}".format(self.driverT, corrId);
+			result = self.db.select(cmd)
+
+		assert (len(result) == 1)
+		info = result[0]
+		start = info['start']
+		end = info['end']
+		pid = info['processId']
+		tid = info['threadId']
+		tid = tid & 0xffffffff	#convert to unsigned
+		assert (end > start)
+		return [start, end, pid, tid]
+
+	def getKernelInfo(self):
+		"""
+		Get GPU kernel info
+		"""
+		cmd = "select name,correlationId,start,end,deviceId,streamId,gridX,gridY,gridZ,blockX,blockY,blockZ from {}".format(self.kernelT)
+		result = self.db.select(cmd)
+		return result
+
+	def getMarkerInfo(self, objId, startTime, endTime):
+		"""
+		This function first finds all NVTX markers encapsulating
+		a runtime / driver kernel launch.
+		It then splits the markers into many lists.
+			layerMarkers : User added NVTX markers
+			traceMarkers : Call trace markers (inserted by pyprof)
+			reprMarkers  : Markers containing the extra_repr() of a module (inserted by pyprof)
+			pyprofMarkers: Markers containing args and kwargs (tensor shape, datatype etc.)
+			seqMarkers   : Markers containing PyTorch internal sequence markers (inserted by PyTorch)
+			altSeqMarkers: Markers inserted by PyTorch between two kernel launches. Needs better explanation.
+			otherMarkers : Markers not in either of the above categories.
+
+		We extract seqId from the seq and altSeq markers. The seqId is used in bprop.
+		We also extract information from the layerMarkers.
+		"""
+
+		layerMarkers = []
+		traceMarkers = []
+		reprMarkers = []
+		pyprofMarkers = []
+		seqMarkers = []
+		otherMarkers = []
+		altSeqMarkers = []
+		bprop = False
+
+		#Helper functions
+
+		def delete(objId, sTime):
+			"""
+			Delete rows from the temporary SQL table which are no longer required.
+			This speeds up future queries.
+			"""
+			margin = 0
+			cmd = 'DELETE FROM marker WHERE objectId = "{}" AND endTime < {}'.format(objId, sTime - margin)
+			#cmd = 'DELETE FROM marker WHERE endTime < {}'.format(sTime - margin)
+			self.db.execute(cmd)
+
+		def getLayerName(mlist):
+			"""
+			Get layer names from layer marker list.
+			"""
+			layers = []
+			assert(type(mlist) == list)
+			for m in mlist:
+				assert("layer:" in m)
+				l = m.split(":")[1]
+				layers.append(l)
+			return layers
+
+		def getSeqId(mlist):
+			"""
+			Get sequence ids from seq / alt seq marker list.
+			"""
+			ids = []
+			assert(type(mlist) == list)
+			for m in mlist:
+				assert(", seq = " in m)
+				seq = int(m.split("=")[1])
+				ids.append(seq)
+
+			#Remove duplicates
+			ids = list(set(ids))
+			ids.sort()
+			return ids
+
+		def seqcompare(elem):
+			"""
+			Sorting function for sequence markers
+			"""
+			assert (", seq = " in elem)
+			#sort by sequence id and then the string
+			l = elem.split(" = ")
+			return l[1] + l[0]
+
+		def prune(mlist):
+			"""
+			Remove markers with the same seqId and if the strings are similar.
+			This function works on a sorted sequence.
+			"""
+			assert (type(mlist) == list)
+			assert (len(mlist))
+			a = mlist[0:1]
+			for i in range(1,len(mlist)):
+				m = mlist[i]
+				pm = mlist[i-1]
+				name,seq = m.split(",")
+				pname,pseq = pm.split(",")
+				similar = (name in pname) or (pname in name)
+				if (seq == pseq) and similar:
+					continue
+				else:
+					a.append(m)
+			return a
+
+		def filterTrace(mlist):
+			"""
+			Filter trace markers to remove certain file names.
+			"""
+			assert (type(mlist) == list)
+			if len(mlist) == 0:
+				return mlist
+			mlist = mlist[-1]	#The last stack trace will be a super set.
+			mlist = eval(mlist)
+			mlist = mlist['traceMarker']
+			assert (type(mlist) == list)
+			mlist = list(filter(lambda x : "/torch/nn/modules/" not in x, mlist))
+			mlist = list(filter(lambda x : "/torch/nn/functional.py" not in x, mlist))
+			mlist = list(filter(lambda x : "/torch/tensor.py" not in x, mlist))
+			mlist = list(filter(lambda x : "/torch/autograd/__init__.py" not in x, mlist))
+			mlist = list(filter(lambda x : "/torch/_jit_internal.py" not in x, mlist))
+			mlist = list(filter(lambda x : "/pyprof/nvtx/nvmarker.py" not in x, mlist))
+			mlist = list(filter(lambda x : "/apex/optimizers/" not in x, mlist))
+			mlist = list(filter(lambda x : "/torch/_utils.py" not in x, mlist))
+			mlist = list(filter(lambda x : "/torch/optim/" not in x, mlist))
+			return mlist
+
+		#Find all encapsulating markers
+		cmd = 'SELECT id,name from marker where \
+				objectId = "{}" and \
+				startTime < {} and \
+				endTime > {} \
+				ORDER BY startTime ASC'.format(objId, startTime, endTime)
+		result = self.db.select(cmd)
+
+		#Bin markers into different lists
+		for r in result:
+			m = self.getString(r['name'])
+
+			#Hack: If its a known gradient checkpointing marker, ignore it.
+			if m.find("CheckpointFunctionBackward") >= 0:
+				continue
+
+			if ("_backward, seq =" in m) or ("Backward, seq =" in m) or ("Backward0, seq =" in m):
+				bprop = True
+
+			if ("mod" in m) and ("op" in m) and ("args" in m) and ("type" in m):
+				pyprofMarkers.append(m)
+			elif ("layer:" in m):
+				layerMarkers.append(m)
+			elif ("traceMarker" in m):
+				traceMarkers.append(m)
+			elif ("strRepr" in m):
+				reprMarkers.append(m)
+			elif (", seq = " in m):
+				seqMarkers.append(m)
+			else:
+				otherMarkers.append(m)
+
+		#Remove duplicates, sort and prune seqMarkers
+		if (len(seqMarkers)):
+			seqMarkers = list(set(seqMarkers))
+			seqMarkers.sort(key=seqcompare)
+			seqMarkers = prune(seqMarkers)
+
+		#Remove duplicates from otherMarkers
+		otherMarkers = list(set(otherMarkers))
+
+		#Get markers with seq id (inserted by PyTorch) from the previous kernel to the present kernel
+		#Only for fprop kernels
+		if (len(result) and not bprop):
+			loId = self.markerId
+			hiId = result[-1]['id']
+			self.markerId = hiId
+			
+			#Get markers between loId and hiId
+			cmd = 'SELECT id,name from marker where objectId = "{}" and id > {} and id < {} ORDER BY startTime ASC'.format(objId, loId, hiId)
+			result1 = self.db.select(cmd)
+
+			for r in result1:
+				m = self.getString(r['name'])
+				#Get only markers with seq id
+				if (", seq=" in m):
+					altSeqMarkers.append(m)
+
+			#Remove duplicates, sort and prune altSeqMarkers
+			if (len(altSeqMarkers)):
+				altSeqMarkers = list(set(altSeqMarkers))
+				altSeqMarkers.sort(key=seqcompare)
+				altSeqMarkers = prune(altSeqMarkers)
+
+		delete(objId, startTime)
+
+		return layerMarkers, filterTrace(traceMarkers), reprMarkers, pyprofMarkers, seqMarkers, otherMarkers, altSeqMarkers, getSeqId(seqMarkers), getSeqId(altSeqMarkers), getLayerName(layerMarkers)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/parse.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/parse.py
new file mode 100644
index 0000000000..c119fe5e09
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/parse/parse.py
@@ -0,0 +1,122 @@
+#!/usr/bin/env python3
+
+"""
+Parse the SQL db and print a dictionary for every kernel.
+"""
+
+import sys
+import argparse
+from tqdm import tqdm
+
+from .db import DB
+from .kernel import Kernel
+from .nvvp import NVVP
+
+def parseArgs():
+	parser = argparse.ArgumentParser(prog=sys.argv[0], description="Parse SQL (nvvp) db.")
+	parser.add_argument("file",
+		type=str,
+		default=None,
+		help="SQL db (nvvp) file.")
+
+	args = parser.parse_args()
+	return args
+
+def main():
+	args = parseArgs()
+
+	db = DB(args.file)
+	nvvp = NVVP(db)
+
+	kInfo = nvvp.getKernelInfo()
+	if len(kInfo) == 0:
+		print("Found 0 kernels. Exiting.", file=sys.stderr)
+		db.close()
+		sys.exit(0)
+	else:
+		print("Found {} kernels. Getting info for each kernel.".format(len(kInfo)), file=sys.stderr)
+
+	nvvp.createMarkerTable()
+
+	prevSeqId = -1
+	prevSubSeqId = -1
+	prevOp = "na"
+
+	Kernel.profStart = nvvp.getProfileStart()
+
+	for i in tqdm(range(len(kInfo)), ascii=True):
+		info = kInfo[i]
+		k = Kernel()
+
+		#Set kernel info
+		k.setKernelInfo(info)
+
+		#Get, set kernel name
+		name = nvvp.getString(k.kNameId)
+		k.setKernelName(name)
+
+		#Get runtime info
+		info = nvvp.getCPUInfo(k.corrId)
+		k.setRunTimeInfo(info)
+
+		#Get and set marker and seqid info
+		info = nvvp.getMarkerInfo(k.objId, k.rStartTime, k.rEndTime)
+		k.setMarkerInfo(info)
+
+		#If the seqId contains both 0 and non zero integers, remove 0.
+		if any(seq != 0 for seq in k.seqId) and (0 in k.seqId):
+			k.seqId.remove(0)
+
+		#Set direction (it uses seq id)
+		k.setDirection()
+
+		#Set op
+		k.setOp()
+
+		#The following code is based on heuristics.
+		#TODO: Refactor.
+		#Assign subSeqId, adjust seqId and altSeqId
+		#seqId can be 0.
+		#A kernel can have multiple seqIds both in fprop and bprop.
+		#In bprop, seqIds might not decrease monotonically. I have observed a few blips.
+		if len(k.seqId):
+			assert (k.dir in ["fprop", "bprop"])
+			if (k.dir == "fprop"):
+				#Check if there is a sequence id larger than the previous
+				inc = (k.seqId[-1] > prevSeqId)
+				if inc:
+					currSeqId = [x for x in k.seqId if x > prevSeqId][0]
+				else:
+					currSeqId = prevSeqId
+			else:
+				currSeqId = k.seqId[0]
+
+			#if ((currSeqId == prevSeqId) and (k.op == prevOp)):
+			if ((currSeqId == prevSeqId) and (k.op == prevOp)) or ((k.op[0] == "forward") and (k.op == prevOp) and (k.mod[0] in ["LSTMCell", "GRUCell", "RNNCell"])):
+				#The second condition is to trap cases when pytorch does not use cudnn for a LSTMCell.
+				k.subSeqId = prevSubSeqId + 1
+
+			prevSeqId = currSeqId
+			prevSubSeqId = k.subSeqId
+			prevOp = k.op
+
+			#Keep currSeqId in k.seqId, move everything else to k.altSeqId
+			for s in k.seqId:
+				if s != currSeqId:
+					k.seqId.remove(s)
+					k.altSeqId.append(s)
+
+			for s in k.altSeqId:
+				if s == currSeqId:
+					k.altSeqId.remove(s)
+
+			k.altSeqId = list(set(k.altSeqId))
+			if (len(k.altSeqId)):
+				(k.altSeqId).sort()
+
+		k.print()
+
+	db.close()
+
+if __name__ == '__main__':
+	main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/__init__.py
new file mode 100644
index 0000000000..cedb1ad415
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/__init__.py
@@ -0,0 +1 @@
+from . import data, prof
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/__main__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/__main__.py
new file mode 100644
index 0000000000..a114f7c948
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/__main__.py
@@ -0,0 +1,10 @@
+import warnings
+
+try:
+    from .prof import main
+except ImportError as e:
+    warnings.warn("Did you make sure to install PyProf dependencies by using the --pyprof flag during Apex installation?")
+    raise e
+
+if __name__ == '__main__':
+    main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/activation.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/activation.py
new file mode 100644
index 0000000000..e2443b0930
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/activation.py
@@ -0,0 +1,65 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class Activation(OperatorLayerBase):
+	"""
+	This class handles the various activation functions.
+	"""
+
+	ops = ["celu", "elu", "elu_", "hardshrink", "hardtanh", "hardtanh_", "leaky_relu", "leaky_relu_", "logsigmoid", "prelu", "relu", "relu_", "relu6", "rrelu", "rrelu_", "selu", "sigmoid", "softplus", "softshrink", "softsign", "tanh", "tanhshrink", "threshold", "threshold_"]
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod in ["torch.nn.functional", "torch", "Tensor"])
+
+		#Filter out named parameters
+		args = list(filter(lambda x : x['name'] == '', args))
+
+		assert (len(args) >= 1)
+		arg = args[0]
+		assert (arg['type'] == "tensor")
+
+		self.i = arg
+		self.dir = d.dir
+
+	def params(self):
+		p = OrderedDict([('T', self.i['shape']),('type', self.i['dtype'])])
+		return p
+
+	def flops(self):
+		direction = self.dir
+		tensor = self.i['shape']
+		t = self.i['dtype']
+
+		# TODO: revise
+		elems = Utility.numElems(tensor)
+		return elems
+
+	def bytes(self):
+		direction = self.dir
+		tensor = self.i['shape']
+		t = self.i['dtype']
+
+		elems = Utility.numElems(tensor)
+		elems = elems * (2 if direction == "fprop" else 3)
+
+		return elems * Utility.typeToBytes(t)
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/base.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/base.py
new file mode 100644
index 0000000000..adfc4ce8a6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/base.py
@@ -0,0 +1,47 @@
+from abc import ABC, abstractmethod
+
+class OperatorLayerBase(ABC):
+	"""
+	Base class for all layers and operators.
+	Every derived class should have the following functions.
+	"""
+
+	@abstractmethod
+	def tc(self):
+		"""
+		Tensor core usage by the kernel.
+		Return "1" (yes), "0" (no, but possible), "-" (not applicable)
+		"""
+		pass
+
+	@abstractmethod
+	def params(self):
+		"""
+		Kernel parameters to be printed.
+		"""
+		pass
+
+	@abstractmethod
+	def flops(self):
+		"""
+		Note that 1 FMA = 2 flops.
+		"""
+		pass
+
+	@abstractmethod
+	def bytes(self):
+		pass
+
+	@abstractmethod
+	def mod(self):
+		"""
+		Name of the module/class e.g. torch.nn.functional.
+		"""
+		pass
+
+	@abstractmethod
+	def op(self):
+		"""
+		Name of the operator e.g. sigmoid.
+		"""
+		pass
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/blas.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/blas.py
new file mode 100644
index 0000000000..706205c7b8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/blas.py
@@ -0,0 +1,340 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+import numpy as np
+
+TC_GEMMS = ["884gemm", "1688gemm"]
+
+class Addmm(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod in ["torch", "Tensor",])
+		assert (op in ["addmm", "addmm_",])
+
+		#Get alpha and beta
+		alpha = 1
+		beta = 1
+		if any(x['name'] == 'alpha' for x in args):
+			alpha = list(filter(lambda x : x['name'] == "alpha", args))[0]
+			alpha = alpha['value']
+
+		if any(x['name'] == 'beta' for x in args):
+			beta = list(filter(lambda x : x['name'] == "beta", args))[0]
+			beta = beta['value']
+
+		self.alpha = alpha
+		self.beta = beta
+
+		#Filter out named parameters
+		args = list(filter(lambda x : x['name'] == '', args))
+
+		assert (len(args) == 3)
+		C,A,B = args
+		m,k1 = A['shape']
+		k2,n = B['shape']
+		assert (k1 == k2)
+		t1 = A['dtype']
+		t2 = B['dtype']
+		t3 = C['dtype']
+		assert(t1 == t2 == t3)
+
+		self.A = A
+		self.B = B
+		self.C = C
+
+		self.m = m
+		self.n = n
+		self.k = k1
+		self.type = t1
+		self.name = d.name
+
+		return
+
+	def tc(self):
+            for s in TC_GEMMS:
+                if s in self.name:
+                    return 1
+            return 0
+
+	def bytes(self):
+		m, n, k = self.m, self.n, self.k
+		return Utility.typeToBytes(self.type) * (m*n + m*k + n*k)
+
+	def flops(self):
+		return self.m * self.n * self.k * 2
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def params(self):
+		p = OrderedDict([('M',self.n),('N',self.m),('K',self.k),('type',self.type)])
+		return p
+
+class Bmm(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "torch") and (op == "bmm")
+
+		#Filter out named params (kwargs)
+		args = list(filter(lambda x : x['name'] == "", args))
+
+		assert (len(args) == 2)
+		A,B = args
+		b1,m,k1 = A['shape']
+		b2,k2,n = B['shape']
+		assert (b1 == b2)
+		assert (k1 == k2)
+		t1 = A['dtype']
+		t2 = B['dtype']
+		assert(t1 == t2)
+
+		self.A = A
+		self.B = B
+		self.b = b1
+		self.m = m
+		self.n = n
+		self.k = k1
+		self.type = t1
+		self.name = d.name
+
+	def tc(self):
+            for s in TC_GEMMS:
+                if s in self.name:
+                    return 1
+            return 0
+
+	def params(self):
+		#p = OrderedDict([('A', A['shape']), ('B', B['shape']), ('type', t1)])
+		p = OrderedDict([('B',self.b), ('M',self.n),('N',self.m),('K',self.k),('type',self.type)])
+		return p
+
+	def flops(self):
+		return self.b * self.m * self.n * self.k * 2
+
+	def bytes(self):
+		b, m, n, k = self.b, self.m, self.n, self.k
+		return Utility.typeToBytes(self.type) * b * (m*n + m*k + n*k)
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+class Matmul(OperatorLayerBase):
+
+	NON_GEMM = ["kernelPointwiseApply2", "reduce_1Block_kernel", "elementwise_kernel"]
+	NON_TC = NON_GEMM + ["dot_kernel"]
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		self.name = d.name
+		self.sub = d.sub
+
+		assert ((mod == "torch") and (op == "matmul")) or ((mod == "Tensor") and (op == "__matmul__"))
+		assert (len(args) == 2)
+
+		assert any([x in d.name for x in Matmul.NON_TC + ["gemm", "gemv"]])
+
+		A,B = args
+		t1 = A['dtype']
+		t2 = B['dtype']
+		assert(t1 == t2)
+
+		A = A['shape']
+		B = B['shape']
+
+		self.A = A
+		self.B = B
+		self.type = t1
+
+		# batch, MNK
+		if (len(A) == 1) and (len(B) == 1):
+			#dot product
+			assert (A[0] == B[0])
+			self.b = (1,)
+			self.m = 1
+			self.n = 1
+			self.k = A[0]
+
+		elif (len(A) == 2) and (len(B) == 2):
+			#gemm
+			m,k1 = A
+			k2,n = B
+			assert(k1 == k2)
+			self.b = (1,)
+			self.m = m
+			self.n = n
+			self.k = k1
+
+		elif (len(A) == 1) and (len(B) == 2):
+			#vector matrix
+			k1 = A[0]
+			k2,n = B
+			assert(k1 == k2)
+
+			self.b = (1,)
+			self.m = 1
+			self.n = n
+			self.k = k1
+
+		elif (len(A) == 2) and (len(B) == 1):
+			#gemv
+			m,k1 = A
+			k2 = B[0]
+			assert (k1 == k2)
+
+			self.b = (1,)
+			self.m = m
+			self.n = 1
+			self.k = k1
+
+		elif (len(A) == 1) and (len(B) > 2):
+			assert (A[0] == B[-2])
+
+			self.b = B[0:-2]
+			self.m = 1
+			self.n = B[-1]
+			self.k = B[-2]
+
+		elif (len(B) == 1) and (len(A) > 2):
+			assert (B[0] == A[-1])
+
+			self.b = A[0:-2]
+			self.m = A[-2]
+			self.n = 1
+			self.k = A[-1]
+
+		else:
+			assert (len(A) >= 2)
+			assert (len(B) >= 2)
+			assert (A[-1] == B[-2])
+			self.m = A[-2]
+			self.n = B[-1]
+			self.k = A[-1]
+
+			aa = np.empty(A[0:-2])
+			bb = np.empty(B[0:-2])
+			self.b = np.broadcast(aa, bb).shape
+
+	def params(self):
+		return OrderedDict([('A', self.A), ('B', self.B), ('type', self.type)])
+
+	def tc(self):
+		if self.name in Matmul.NON_TC:
+			return "-"
+		else:
+                    for s in TC_GEMMS:
+                        if s in self.name:
+                            return 1
+                    return 0
+
+	def bytes(self):
+		# TODO: check bytes for non-GEMM cases
+		if self.name in Matmul.NON_GEMM:
+			return 2 * Utility.typeToBytes(self.type) * Utility.numElems(self.A) #could be B as well
+		else:
+			m, n, k = self.m, self.n, self.k
+			return Utility.typeToBytes(self.type) * (m*n + m*k + n*k)
+
+	def flops(self):
+		# TODO: calculate actual FLOPs. At least we're not saying it's GEMM FLOPs for now.
+		if self.name in Matmul.NON_GEMM:
+			return 0
+		else:
+			return Utility.numElems(self.b) * self.m * self.n * self.k * 2
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+class Mm(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "torch") and (op == "mm")
+		assert (len(args) == 2)
+
+		A,B = args
+		m,k1 = A['shape']
+		k2,n = B['shape']
+		assert (k1 == k2)
+		t1 = A['dtype']
+		t2 = B['dtype']
+		assert(t1 == t2)
+
+		self.A = A
+		self.B = B
+		self.m = m
+		self.n = n
+		self.k = k1
+		self.type = t1
+		self.name = d.name
+
+		return
+
+	def params(self):
+		p = OrderedDict([('M',self.n),('N',self.m),('K',self.k),('type',self.type)])
+		return p
+
+	def tc(self):
+            for s in TC_GEMMS:
+                if s in self.name:
+                    return 1
+            return 0
+
+	def bytes(self):
+		m, n, k = self.m, self.n, self.k
+		return Utility.typeToBytes(self.type) * (m*n + m*k + n*k)
+
+	def flops(self):
+		return self.m * self.n * self.k * 2
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/conv.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/conv.py
new file mode 100644
index 0000000000..df388df133
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/conv.py
@@ -0,0 +1,236 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class Conv(OperatorLayerBase):
+
+	"""
+	# N = batch size
+	# C,H,W = input channels, height, width
+	# K,P,Q = output channels, height, width
+	# R,S = filter height, width
+	# g = groups
+	"""
+
+	#todo: refine winograd and FFT
+	convAuxList = ["nchwToNhwc", "nhwcToNchw", "OffsetsKernel",]
+	winoAuxList = ["generateWinogradTilesKernel", "winogradWgradData", "winogradWgradOutput", "winogradWgradDelta"]
+	fftAuxList = ["compute_gemm_pointers", "flip_filter", "fft2d_r2c_", "fft2d_c2r_", "fft1d_r2c", "fft1d_c2r"]
+	miscAuxList = ["scaleTensor_kernel",]
+
+	convList = ["_s884cudnn_", "_s1688cudnn_", "_scudnn_", "2d_grouped_direct_kernel", "cudnn::detail::implicit_convolve_sgemm", "cudnn::detail::dgrad2d_alg1_1", "cudnn::detail::wgrad_alg0_engine", "cudnn::detail::dgrad_engine", "dgrad_1x1_stride_2x2", "spatialDepthwiseConvolutionUpdateOutput"]
+	winoList = ["winograd3x3Kernel", "_sgemm_"]
+	fftList = ["fermiPlusCgemmLDS128_batched", "_gcgemm_",]
+	miscList = []
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		self.dir = d.dir
+		self.name = d.name
+		self.sub = d.sub
+
+		assert (mod == "torch.nn.functional")
+		assert (op in ["conv1d", "conv2d"])
+		length = len(args)
+		assert (length >= 2) and (length <= 7)
+		i,w = args[0], args[1]
+		assert (i['type'] == "tensor")
+		assert (w['type'] == "tensor")
+
+		#ignore bias
+
+		if (length >= 4) and (args[3]['name'] == ""):
+			s = args[3]
+		elif any(x['name'] == 'stride' for x in args):
+			s = list(filter(lambda x : x['name'] == 'stride', args))[0]
+		else:
+			s = {'name': 'stride', 'type': 'int', 'value': 1}
+
+		if (length >= 5) and (args[4]['name'] == ""):
+			p = args[4]
+		elif any(x['name'] == 'padding' for x in args):
+			p = list(filter(lambda x : x['name'] == 'padding', args))[0]
+		else:
+			p = {'name': 'padding', 'type': 'int', 'value': 0}
+
+		if (length >= 6) and (args[5]['name'] == ""):
+			d = args[5]
+		elif any(x['name'] == 'dilation' for x in args):
+			d = list(filter(lambda x : x['name'] == 'dilation', args))[0]
+		else:
+			d = {'name': 'dilation', 'type': 'int', 'value': 1}
+
+		if (length == 7) and (args[6]['name'] == ""):
+			g = args[6]
+		elif any(x['name'] == 'groups' for x in args):
+			g = list(filter(lambda x : x['name'] == 'groups', args))[0]
+		else:
+			g = {'name': 'groups', 'type': 'int', 'value': 1}
+
+		if op == "conv1d":
+			assert (len(i['shape']) == 3)
+			assert (len(w['shape']) == 3)
+			assert (i['dtype'] == w['dtype'])
+			N, C1, W = i['shape']
+			K, C2, S = w['shape']
+			assert (C1 == C2)
+			p = p['value'] if Utility.isscalar(p['type']) else p['value'][0]
+			s = s['value'] if Utility.isscalar(s['type']) else s['value'][0]
+			d = d['value'] if Utility.isscalar(d['type']) else d['value'][0]
+			g = g['value']
+			assert (g == 1)
+			H = 1
+			R = 1
+
+			P = 1 + (H - (((R-1))+1))
+			Q = 1 + (W + 2*p - (((S-1)*d)+1))/s
+			P = int(P)
+			Q = int(Q)
+			if (H == 1):
+				assert (P == 1)
+			if (W == 1):
+				assert (Q == 1)
+
+			self.N = N
+			self.C = C1
+			self.H = H
+			self.W = W
+			self.K = K
+			self.P = P
+			self.Q = Q
+			self.R = R
+			self.S = S
+			self.ph = 0
+			self.pw = p
+			self.U = 1
+			self.V = s
+			self.dh = 1
+			self.dw = d
+			self.g = g
+			self.type = i['dtype']
+
+		elif op == "conv2d":
+			assert (len(i['shape']) == 4)
+			assert (len(w['shape']) == 4)
+			assert (i['dtype'] == w['dtype'])
+			N, C1, H, W = i['shape']
+			K, C2, R, S = w['shape']
+
+			if Utility.isscalar(p['type']):
+				ph = pw = p['value']
+			else:
+				assert (p['type'] == "tuple")
+				ph, pw = p['value']
+
+			if Utility.isscalar(s['type']):
+				sh = sw = s['value']
+			else:
+				assert (s['type'] == "tuple")
+				sh, sw = s['value']
+
+			if Utility.isscalar(d['type']):
+				dh = dw = d['value']
+			else:
+				assert (d['type'] == "tuple")
+				dh, dw = d['value']
+
+			g = g['value']
+			assert (g >= 1)
+			assert (C1 == C2*g)
+
+			P = 1 + (H + 2*ph - (((R-1)*dh)+1))/sh
+			Q = 1 + (W + 2*pw - (((S-1)*dw)+1))/sw
+			P = int(P)
+			Q = int(Q)
+			if (H == 1):
+				assert (P == 1)
+			if (W == 1):
+				assert (Q == 1)
+
+			self.N = N
+			self.C = C1
+			self.H = H
+			self.W = W
+			self.K = K
+			self.P = P
+			self.Q = Q
+			self.R = R
+			self.S = S
+			self.ph = ph
+			self.pw = pw
+			self.U = sh
+			self.V = sw
+			self.dh = dh
+			self.dw = dw
+			self.g = g
+			self.type = i['dtype']
+
+		else:
+			assert False
+
+	def params(self):
+		p = OrderedDict([('N',self.N), ('C',self.C), ('H',self.H), ('W',self.W), ('K',self.K), ('P',self.P), ('Q',self.Q), ('R',self.R), ('S',self.S), ('ph',self.ph), ('pw',self.pw), ('U',self.U), ('V',self.V), ('dh',self.dh), ('dw',self.dw), ('g',self.g), ('type',self.type)])
+		return p
+
+	def conv_bytes_flops(self, N, C, H, W, K, P, Q, R, S, g, t):
+		f = 2*N*K*P*Q*C*R*S/g #for fprop
+		elems = N*C*H*W + K*C*R*S/g + N*K*P*Q
+		b = elems * Utility.typeToBytes(t)
+		return b,f
+
+	def bytes_flops(self):
+		N,C,H,W,K,P,Q,R,S,ph,pw,U,V,dh,dw,g,t = self.params().values()
+
+		if any(x in self.name for x in Conv.convAuxList+Conv.winoAuxList+Conv.fftAuxList+Conv.miscAuxList):
+			bytes, flops = [0, 0]
+
+		elif any(x in self.name for x in Conv.convList+Conv.winoList+Conv.fftList+Conv.miscList):
+			if g == 1:
+				bytes, flops = self.conv_bytes_flops(N,C,H,W,K,P,Q,R,S,g,t)
+			else:
+				if "2d_grouped_direct_kernel" in self.name:	#only 1 kernel is called
+					bytes, flops = self.conv_bytes_flops(N,C,H,W,K,P,Q,R,S,g,t)
+				elif "spatialDepthwiseConvolutionUpdateOutput" in self.name: #one kernel for separable conv
+					bytes, flops = self.conv_bytes_flops(N,C,H,W,K,P,Q,R,S,g,t)
+				else:	#a kernel per group is called
+					bytes, flops = self.conv_bytes_flops(N,C/g,H,W,K/g,P,Q,R,S,1,t)
+
+		elif ("calc_bias_diff" in self.name):	#bias gradient
+			elems = N*K*P*Q
+			flops = elems
+			bytes = 2 * elems * Utility.typeToBytes(t)
+			#params = OrderedDict([('N',N), ('K',K), ('P',P), ('Q',Q), ('type', t)])
+
+		else:
+			bytes, flops = [0, 0]
+
+		return bytes, flops
+
+	def bytes(self):
+		b,_ = self.bytes_flops()
+		return b
+
+	def flops(self):
+		_,f = self.bytes_flops()
+		return f
+
+	def tc(self):
+		for s in ["884cudnn", "1688cudnn"]:
+			if s in self.name:
+				return 1
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/convert.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/convert.py
new file mode 100644
index 0000000000..0d6735a819
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/convert.py
@@ -0,0 +1,62 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class Convert(OperatorLayerBase):
+	"""
+	Class to handle convert operations.
+	"""
+	ops = ["byte", "char", "double", "float", "half", "int", "long", "short", "to"]
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "Tensor")
+		assert (op in Convert.ops)
+		assert (len(args) == 1)
+
+		#The argument could be a tensor or scalar
+		t = args[0]
+		if t['type'] == "tensor":
+			shape = t['shape']
+			stype = t['dtype']
+		else:
+			shape = (1,)
+			stype = t['type']
+		if self.op_ == "to":
+			op = stype
+
+		self.shape = shape
+		self.stype = stype
+		self.dtype = op
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('stype', self.stype), ('dtype', self.dtype)])
+		return p
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def tc(self):
+		return "-"
+
+	def elems(self):
+		return Utility.numElems(self.shape)
+
+	def flops(self):
+		return 0
+
+	def bytes(self):
+		b = self.elems() * (Utility.typeToBytes(self.stype) + Utility.typeToBytes(self.dtype))
+		return b
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/data.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/data.py
new file mode 100644
index 0000000000..dcd323dea6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/data.py
@@ -0,0 +1,54 @@
+from .utility import Utility
+
+class Data(object):
+	"""
+	Class to store all the data for every kernel e.g. name, bytes, flops, device, stream etc.
+	"""
+	def __init__(self, kernel):
+		#Available from NVprof
+		self.tid = kernel['tid']
+		self.device = kernel['device']
+		self.stream = kernel['stream']
+		self.grid = str(kernel['grid']).replace(" ","").replace("(","").replace(")","")
+		self.block = str(kernel['block']).replace(" ","").replace("(","").replace(")","")
+		self.name = kernel['kShortName'].replace(" ","_")
+		self.lName = kernel['kLongName']
+		self.sil = kernel['kDuration']	#units ns
+
+		self.index = None
+
+		#Markers
+		self.argMarker = kernel['marker']
+		self.modMarker = kernel['reprMarkers']
+		self.seqMarker = kernel['seqMarker']
+
+		self.layer = kernel['layer']
+		self.trace = kernel['trace']
+
+		self.seqId = kernel['seqId']
+		self.altSeqId = kernel['altSeqId']
+
+		self.dir = kernel['dir']
+		self.sub = kernel['subSeqId']
+
+		self.mod = "na"
+		self.op = "na"
+		self.params = {"na":"na"}
+		self.tc = "na"
+		self.flops = 0
+		self.bytes = 0
+
+	def setParams(self, params):
+		#Remove space from params
+		qaz = ""
+		for key,value in params.items():
+			if "type" not in key:
+				qaz += "{}={},".format(key,value)
+			else:
+				if type(value) is str:
+					qaz += "{},".format(Utility.typeToString(value))
+				else:
+					qaz += "{}".format(value)
+
+		self.params = qaz.replace(" ", "")
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/dropout.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/dropout.py
new file mode 100644
index 0000000000..47b352bbf1
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/dropout.py
@@ -0,0 +1,50 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class Dropout(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "torch.nn.functional")
+		assert (op == "dropout")
+		#assert (len(args) == 1)
+
+		self.shape = args[0]['shape']
+		self.type  = args[0]['dtype']
+		self.dir = d.dir
+
+		return
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def tc(self):
+		return "-"
+
+	def elems(self):
+		return Utility.numElems(self.shape)
+
+	def bytes(self):
+		#Ignoring the cost of writing and reading the mask
+		return Utility.typeToBytes(self.type) * self.elems() * 2
+
+	def flops(self):
+		# Note: This is approximate and depends on the RNG
+		return 5*self.elems()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/embedding.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/embedding.py
new file mode 100644
index 0000000000..32a1c35eaa
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/embedding.py
@@ -0,0 +1,71 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class Embedding(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "torch.nn.functional")
+		assert (op == "embedding")
+
+		self.ishape = args[0]['shape']
+		self.itype = args[0]['dtype']
+
+		self.eshape = args[1]['shape']
+		self.etype = args[1]['dtype']
+
+		assert (len(self.eshape) == 2)
+
+		self.dir = d.dir
+		self.sub = d.sub
+		return
+
+	def params(self):
+		p = OrderedDict([('I', self.ishape), ('itype', self.itype), ('E', self.eshape), ('etype', self.etype)])
+		return p
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def tc(self):
+		return "-"
+
+	def bytes(self):
+		ishape = self.ishape
+		itype = self.itype
+		eshape = self.eshape
+		etype = self.etype
+
+		ielems = Utility.numElems(ishape)
+
+		b = 0
+		if self.dir == "fprop":
+			#indices
+			b += ielems * Utility.typeToBytes(itype)
+			#read and write the embedding matrix
+			b += ielems * eshape[1] * 2 * Utility.typeToBytes(etype)
+		else:
+			#3 times the size of the incoming gradient
+			b = ielems * eshape[1] * 3 * Utility.typeToBytes(etype)
+
+			if self.sub > 0:
+				b = 0
+
+		return b
+
+	def flops(self):
+		# Note: not implemented yet
+		return 0
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/index_slice_join_mutate.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/index_slice_join_mutate.py
new file mode 100644
index 0000000000..1ecbe60d73
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/index_slice_join_mutate.py
@@ -0,0 +1,419 @@
+from collections import OrderedDict
+from .utility import Utility
+import numpy as np
+from .base import OperatorLayerBase
+
+class Cat(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "torch")
+		assert (op == "cat")
+		assert (len(args) >= 2)
+
+		t = args[0]['dtype']
+		shapes = []
+
+		for arg in args:
+			if arg['type'] == "tensor":
+				assert (arg['dtype'] == t)
+				shapes.append(arg['shape'])
+
+		self.type = t
+		self.shapes = shapes
+
+	def params(self):
+		p = OrderedDict([('T', self.shapes), ('type', self.type)])
+		return p
+
+	def flops(self):
+		return 0
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def bytes(self):
+		b = 0
+		for s in self.shapes:
+			b += Utility.numElems(s)
+		return 2 * b * Utility.typeToBytes(self.type)
+
+class Reshape(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "Tensor")
+		assert (op == "reshape")
+
+		#Temporarily commenting three lines
+		#assert (len(args) == 2)
+		#t,s = args
+		#assert s['type'] == "tuple"
+
+		t = args[0]
+		assert t['type'] == "tensor"
+		self.type = t['dtype']
+		self.shape = t['shape']
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def flops(self):
+		return 0
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def bytes(self):
+		return 0
+
+class Gather(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "Tensor") or (mod == "torch")
+		assert (op == "gather")
+
+		#Filter out the "out" parameter
+		args = list(filter(lambda x : x['name'] != 'out', args))
+		assert (len(args) == 3)
+
+		#Get input
+		if (args[0]['name'] == ""):
+			arg = args[0]
+		else:
+			arg = list(filter(lambda x : x['name'] == "input", args))[0]
+
+		assert (arg['type'] == "tensor")
+
+		self.shape = arg['shape']
+		self.type = arg['dtype']
+
+	def params(self):
+		p = OrderedDict([('T', self.shape),('type', self.type)])
+		return p
+
+	def flops(self):
+		return 0
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def bytes(self):
+		return 2 * Utility.numElems(self.shape) * Utility.typeToBytes(self.type)
+
+class MaskedScatter(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "Tensor")
+		assert (op == "masked_scatter_")
+		assert (len(args) == 3)
+
+		dst, mask, src = args
+		assert (dst['type'] == mask['type'] == src['type'] == "tensor")
+		assert (mask['dtype'] == "uint8")
+		assert (dst['dtype'] == src['dtype'])
+		assert (dst['shape'] == mask['shape'])
+
+		self.shape = dst['shape']
+		self.type = dst['dtype']
+		self.seqId = d.seqId
+
+	def params(self):
+		p = OrderedDict([('T', self.shape),('type', self.type)])
+		return p
+
+	def flops(self):
+		return 0
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def bytes(self):
+		elems = Utility.numElems(self.shape)
+
+		#src and dst
+		b = 2 * elems * Utility.typeToBytes(self.type)
+
+		#mask (uint8)
+		b += elems
+
+		if (self.seqId > 0):
+			b = 0
+		return b
+
+class Nonzero(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod in ["torch", "Tensor"])
+		assert (op == "nonzero")
+		assert (len(args) == 1)
+
+		arg = args[0]
+		self.shape = arg['shape']
+		self.type = arg['dtype']
+		self.seqId = d.seqId
+
+	def params(self):
+		p = OrderedDict([('T', self.shape),('type', self.type)])
+		return p
+
+	def flops(self):
+		return 0
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def bytes(self):
+		elems = Utility.numElems(self.shape)
+		dim = len(self.shape)
+
+		#input tensor
+		b = elems * Utility.typeToBytes(self.type)
+
+		#in the worst case, the output is a (elems x dim) tensor of type "long"
+		b += elems * dim * Utility.typeToBytes("int64")
+
+		if self.seqId > 0:
+			return 0
+		else:
+			return b
+
+class IndexSelect(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "Tensor") or (mod == "torch")
+		assert (op == "index_select")
+
+		#Filter out the "out" parameter
+		args = list(filter(lambda x : x['name'] != 'out', args))
+		assert (len(args) == 3)
+
+		#Get input, dim and index
+		if (args[0]['name'] == ""):
+			t = args[0]
+		else:
+			t = list(filter(lambda x : x['name'] == "input", args))[0]
+
+		if (args[1]['name'] == ""):
+			d = args[1]
+		else:
+			d = list(filter(lambda x : x['name'] == "dim", args))[0]
+
+		if (args[2]['name'] == ""):
+			i = args[2]
+		else:
+			i = list(filter(lambda x : x['name'] == "index", args))[0]
+
+		assert (t['type'] == i['type'] == "tensor")
+		assert (d['type'] == "int")
+		assert (i['dtype'] == "int64")
+		assert (len(i['shape']) == 1)
+
+		shape = t['shape']
+		dim = d['value']
+		indices = i['shape'][0]
+		assert (dim < len(shape))
+
+		self.shape = shape
+		self.dim = dim
+		self.indices = indices
+		self.type = t['dtype']
+
+	def params(self):
+		p = OrderedDict([('T', self.shape),('D', self.dim),('I', self.indices),('type', self.type)])
+		return p
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def flops(self):
+		return 0
+
+	def bytes(self):
+		#determine the shape of the output tensor
+		shape = list(self.shape)
+		shape[self.dim] = self.indices
+
+		b = 0
+
+		#time to read the input and write the output
+		elems = Utility.numElems(shape)
+		b += 2 * elems * Utility.typeToBytes(self.type)
+
+		#time to read the indices
+		b += self.indices * Utility.typeToBytes("int64")
+
+		return b
+
+class MaskedSelect(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+		self.sub = d.sub
+
+		assert (mod == "Tensor") or (mod == "torch")
+		assert (op == "masked_select")
+
+		#Filter out the "out" parameter
+		args = list(filter(lambda x : x['name'] != 'out', args))
+		assert (len(args) == 2)
+
+		#Get input and mask
+		if (args[0]['name'] == ""):
+			t = args[0]
+		else:
+			t = list(filter(lambda x : x['name'] == "input", args))[0]
+
+		if (args[1]['name'] == ""):
+			m = args[1]
+		else:
+			m = list(filter(lambda x : x['name'] == "mask", args))[0]
+
+		assert (m['dtype'] == "uint8")
+
+		tensor = t['shape']
+		mask = m['shape']
+
+		#check for broadcast condition
+		if (tensor != mask):
+			array1 = np.empty(list(tensor))
+			array2 = np.empty(list(mask))
+			try:
+				out = np.broadcast(array1, array2).shape
+			except:
+				assert False
+
+		self.tshape = tensor
+		self.mshape = mask
+		self.type = t['dtype']
+
+	def params(self):
+		p = OrderedDict([('T', self.tshape),('M', self.mshape),('type', self.type)])
+		return p
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def bytes(self):
+		tensor = self.tshape
+		mask = self.mshape
+		t = self.type
+
+		#in the worst case, #output elements = #input elements
+		b = 2 * Utility.numElems(tensor) * Utility.typeToBytes(t)
+
+		#mask tensor (assuming uint8)
+		b += Utility.numElems(mask)
+		return b
+
+	def flops(self):
+		return 0
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/linear.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/linear.py
new file mode 100644
index 0000000000..7c9e13cbef
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/linear.py
@@ -0,0 +1,188 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class Linear(OperatorLayerBase):
+
+	'''
+	Notes:
+	If the bias occurs before the GEMM, then its 1 write (bias expansion).
+	If the bias occurs after, then its 1 read and 1 write.
+	bias in bprop is a reduction and hence is 1 read.
+	'''
+
+	gemmKernels = ["gemm", "gemv", "dot_kernel", "splitKreduce_kernel", "reduce_1Block_kernel"]
+	biasKernels = ["kernelReduceContigDim", "kernelReduceNoncontigDim_shared", "elementwise_kernel", "reduce_kernel"]
+
+	def setXWBMNK(self, args):
+		x = None
+		w = None
+		b = None
+		if (len(args) == 2):
+			x,w = args
+		elif (len(args) == 3):
+			x,w,b = args
+			assert (x['type'] == w['type'] == "tensor")
+			if (b['type'] == "tensor"):
+				assert(len(b['shape']) == 1)
+			elif (b['type'] == "NoneType"):
+				assert b['value'] is None
+				b = None
+			else:
+				assert False
+		else:
+			assert False
+
+		assert(len(w['shape']) == 2)
+		k1 = x['shape'][-1]
+		n,k2 = w['shape']
+		assert(k1 == k2)
+		if b is not None:
+			assert(b['shape'][0] == n)
+		t1 = x['dtype']
+		t2 = w['dtype']
+		assert(t1 == t2)
+
+		# X, W, B
+		self.x = x['shape']
+		self.w = w['shape']
+		self.b = b['shape'] if b is not None else None
+		self.type = t1
+
+		# M, N, K
+		#n = Utility.numElems(x[0:-1])
+		n = self.x[0:-1]
+		k = self.x[-1]
+		m,k1 = self.w
+		assert (k == k1)
+
+		self.m = m
+		self.n = n
+		self.k = k
+
+	def tc(self):
+		if self.op() == "linear":
+			return 1 if "884gemm" in self.name else 0
+		else:
+			return "-"
+
+	def __init__(self, d):
+		self.name = d.name
+		self.dir = d.dir
+		self.sub = d.sub
+
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		assert (mod == "torch.nn.functional")
+		assert (op == "linear")
+
+		self.setXWBMNK(args)
+
+		if any(x in d.name for x in Linear.gemmKernels):
+			self.op_ = "linear"
+		else:
+			assert (d.name in Linear.biasKernels)
+			self.op_ = "bias"
+
+		'''
+		elif (("kernelPointwiseApply2" in d.name) or ("kernelReduceContigDim" in d.name) or ("kernelReduceNoncontigDim_shared" in d.name)):
+			#bias expansion was before the gemm
+			self.op_ = "bias"
+
+		elif ("elementwise_kernel" in d.name):
+			#Bias addition happens later with a broadcast tensor
+			self.op_ = "bias"
+			assert (len(d.argMarker) == 2)
+			marker = eval(d.argMarker[1])
+			mod = marker['mod']
+			op = marker['op']
+			args = marker['args']
+
+			assert (mod == "Tensor")
+			assert (op == "__iadd__")
+			assert (len(args) == 2)
+			mn = args[0]['shape']
+			b = args[1]['shape']
+			assert (len(b) == 1)
+
+			assert (mn == (self.n + (self.m,)))
+			assert (b == self.b)
+
+		else:
+			assert False
+		'''
+
+	def params(self):
+		#p = OrderedDict([('X', self.x), ('W', self.w), ('B', self.b), ('type', self.type)])
+
+		m, n, k, x, w, t = self.m, self.n, self.k, self.x, self.w, self.type
+		if len(n) == 1:
+			n = n[0]
+
+		if self.op_ == "linear":
+			if self.dir == "fprop":
+				p = OrderedDict([('M', m), ('N', n), ('K', k), ('type', t)])
+			elif self.dir == "bprop":
+				if self.sub == 0:		#dgrad (most likely)
+					p = OrderedDict([('M', k), ('N', n), ('K', m), ('type', t)])
+				elif self.sub == 1:	#wgrad (most likely)
+					p = OrderedDict([('M', k), ('N', m), ('K', n), ('type', t)])
+				else:
+					#This happens when there are additional kernels for reduction
+					p = OrderedDict([('X', x), ('W', w), ('type', t)])
+			else:
+				assert False
+
+		elif self.op_ == "bias":
+			p = OrderedDict([('M', m), ('N', n), ('type', t)])
+		else:
+			assert False
+		return p
+
+	def op(self):
+		return self.op_
+
+	def bytesFlops(self):
+
+		m = self.m
+		n = Utility.numElems(self.n)
+		k = self.k
+
+		if self.op_ == "linear":
+			if self.dir == "fprop":
+				f = m * n * k * 2
+				b = m*n + m*k + n*k * Utility.typeToBytes(self.type)
+			elif self.dir == "bprop":
+				if self.sub == 0:		#dgrad (most likely)
+					f = m * n * k * 2
+					b = m*n + m*k + n*k * Utility.typeToBytes(self.type)
+				elif self.sub == 1:	#wgrad (most likely)
+					f = m * n * k * 2
+					b = m*n + m*k + n*k * Utility.typeToBytes(self.type)
+				else:
+					#This happens when there are additional kernels for reduction
+					f = 0
+					b = 0
+			else:
+				assert False
+
+		elif self.op_ == "bias":
+			f = m * n
+			b = 2 * m * n * Utility.typeToBytes(self.type)
+		else:
+			assert False
+		return b,f
+
+	def bytes(self):
+		b, f = self.bytesFlops()
+		return b
+
+	def flops(self):
+		b, f = self.bytesFlops()
+		return f
+
+	def mod(self):
+		return self.mod_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/loss.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/loss.py
new file mode 100644
index 0000000000..3fe8a5501a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/loss.py
@@ -0,0 +1,84 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+#TODO: Add support for additional loss functions.
+
+class MSELoss(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "torch.nn.functional")
+		assert (op == "mse_loss")
+		assert (len(args) == 3)
+
+		#Get input, target and reduction
+		if (args[0]['name'] == ""):
+			x = args[0]
+		else:
+			x = list(filter(lambda x : x['name'] == "input", args))[0]
+
+		if (args[1]['name'] == ""):
+			y = args[1]
+		else:
+			y = list(filter(lambda x : x['name'] == "target", args))[0]
+
+		if (args[2]['name'] == ""):
+			r = args[2]
+		else:
+			r = list(filter(lambda x : x['name'] == "reduction", args))[0]
+
+		assert (x['type'] == y['type'] == "tensor")
+		assert (x['shape'] == y['shape'])
+		assert (x['dtype'] == y['dtype'])
+		assert (r['type'] == "str")
+		assert (r['value'] in ["none", "mean", "sum"])
+
+		self.shape = x['shape']
+		self.type = x['dtype']
+		self.red = r['value']
+		self.dir = d.dir
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type), ('red', self.red)])
+		return p
+
+	def elems(self):
+		red = self.red
+		e = Utility.numElems(self.shape)
+
+		if self.dir == "fprop":
+			if red == "none":
+				e *= 3
+			else:
+				e *= 2
+		else:
+			if red == "none":
+				e *= 4
+			else:
+				e *= 3
+		return e
+
+	def bytes(self):
+		return self.elems() * Utility.typeToBytes(self.type)
+
+	def flops(self):
+		return self.elems() * 2 + 1
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/misc.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/misc.py
new file mode 100644
index 0000000000..e1c247d462
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/misc.py
@@ -0,0 +1,219 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class Foo(OperatorLayerBase):
+	"""
+	An object of Foo is instantiated when we detect an unsupported operator.
+	"""
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		shapes = []
+		types = []
+
+		for arg in args:
+			if arg['type'] == "tensor":
+				shapes.append(arg['shape'])
+				types.append(arg['dtype'])
+
+		self.shape = shapes
+		self.type = types
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def flops(self):
+		return 0
+
+	def bytes(self):
+		return 0
+
+class Copy(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "Tensor")
+		assert (op == "copy_")
+		assert (len(args) == 2)
+
+		dst, src = args
+		assert (src['type'] == dst['type'])
+		assert (src['shape'] == dst['shape'])
+
+		self.shape = src['shape']
+		self.stype = src['dtype']
+		self.dtype = dst['dtype']
+
+	def params(self):
+		#The data type might be different
+		p = OrderedDict([('T', self.shape), ('stype', self.stype), ('dtype', self.dtype)])
+		return p
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def flops(self):
+		return 0
+
+	def elems(self):
+		return Utility.numElems(self.shape)
+
+	def bytes(self):
+		return self.elems() * (Utility.typeToBytes(self.stype) + Utility.typeToBytes(self.dtype))
+
+class Clone(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "Tensor")
+		assert (op == "clone")
+		assert (len(args) == 1)
+		t = args[0]
+		self.shape = t['shape']
+		self.type = t['dtype']
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def flops(self):
+		return 0
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def elems(self):
+		return Utility.numElems(self.shape)
+
+	def bytes(self):
+		return 2 * self.elems() * Utility.typeToBytes(self.type)
+
+class Contiguous(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "Tensor")
+		assert (op == "contiguous")
+		assert (len(args) == 1)
+		t = args[0]
+		self.shape = t['shape']
+		self.type = t['dtype']
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def flops(self):
+		return 0
+
+	def bytes(self):
+		return 2 * Utility.numElems(self.shape) * Utility.typeToBytes(self.type)
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+class Any(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "Tensor")
+		assert (op == "any")
+		assert (len(args) == 1)	#could be 2 as well, the second argument is a bool
+		t = args[0]
+
+		self.shape = t['shape']
+		self.type = t['dtype']
+		self.sub = d.sub
+		return
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def tc(self):
+		return "-"
+
+	def flops(self):
+		return 0
+
+	def bytes(self):
+		return Utility.numElems(self.shape) * Utility.typeToBytes(self.type)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/normalization.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/normalization.py
new file mode 100644
index 0000000000..c9c5ae0b17
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/normalization.py
@@ -0,0 +1,54 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class BatchNorm(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (op == "batch_norm")
+		assert (len(args) == 8)
+		i = args[0]
+		assert (i['type'] == "tensor")
+
+		self.shape = i['shape']
+		self.type = i['dtype']
+		self.dir = d.dir
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def elems(self):
+		return Utility.numElems(self.shape)
+
+	def flops(self):
+		# Variance algo-dependent, but this is a reasonable value.
+		return self.elems() * 8
+
+	def bytes(self):
+		e = self.elems()
+		if self.dir == "fprop":
+			e *= 4
+		else:
+			e *= 5
+
+		return e * Utility.typeToBytes(self.type)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/optim.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/optim.py
new file mode 100644
index 0000000000..f2c32759b4
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/optim.py
@@ -0,0 +1,65 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+#TODO: Add support for other optimizers.
+
+class Adam(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert(op == "adam")
+		assert (len(args) == 12) or (len(args) == 14)
+		w, hw, m, v, g = args[0:5]
+		assert (w['shape'] == m['shape'] == v['shape'] == g['shape'])
+		assert (hw['shape'] == w['shape']) or (hw['shape'] == (0,))		#hw could be null
+		assert (w['type'] == m['type'] == v['type'] == g['type'] == hw['type'] == "tensor")
+		assert (w['dtype'] == m['dtype'] == v['dtype'] == "float32")
+
+		self.w = w
+		self.g = g
+
+	def params(self):
+		p = OrderedDict([('T',self.w['shape']), ('wtype',self.w['dtype']), ('gtype',self.g['dtype'])])
+		return p
+
+	def flops(self):
+		return 0
+
+	def bytes(self):
+		wshape = self.w['shape']
+		wtype = self.w['dtype']
+		gtype = self.g['dtype']
+		b = 0
+
+		elems = Utility.numElems(wshape)
+
+		#Get time to stream read/write w, m, v
+		b += 6 * elems *  Utility.typeToBytes(wtype)
+
+		#Get time to read "g"
+		b += elems * Utility.typeToBytes(gtype)
+
+		if wtype != gtype: #mixed precision
+			#Get time to write "hw
+			b += elems * Utility.typeToBytes(gtype)
+
+		return b
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/output.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/output.py
new file mode 100644
index 0000000000..832514cf24
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/output.py
@@ -0,0 +1,149 @@
+import errno, os, sys
+
+class Output():
+	"""
+	This class handles printing of a columed output and a CSV.
+	"""
+
+	# The table below is organized as 
+	# user_option: [output_header, attribute_in_Data_class, type, min_width_in_columed_output]
+	table = {
+		"idx":		["Idx",			"index",	int,	7],
+		"seq":		["SeqId",		"seqId",	str,	7],
+		"altseq":	["AltSeqId",	"altSeqId",	str,	7],
+		"tid":		["TId",			"tid",		int,	12],
+		"layer":	["Layer", 		"layer",	str,	10],
+		"trace":	["Trace",		"trace",	str,	25],
+		"dir":		["Direction",	"dir",		str,	5],
+		"sub":		["Sub",			"sub",		int,	3],
+		"mod":		["Module",		"mod",		str,	15],
+		"op":		["Op",			"op",		str,	15],
+		"kernel":	["Kernel",		"name",		str,	0],
+		"params":	["Params",		"params",	str,	0],
+		"sil":		["Sil(ns)",		"sil",		int,	10],
+		"tc":		["TC",			"tc",		str,	2],
+		"device":	["Device",		"device",	int,	3],
+		"stream":	["Stream",		"stream",	int,	3],
+		"grid":		["Grid",		"grid",		str,	12],
+		"block":	["Block",		"block",	str,	12],
+		"flops":	["FLOPs", 		"flops",	int,	12],
+		"bytes":	["Bytes",		"bytes", 	int,	12]
+	}
+
+	def __init__(self, args):
+		self.cols = args.c
+		self.csv = args.csv
+		self.col = True if (args.w > 0) else False
+		self.width = args.w
+
+		w = 0
+		for col in self.cols:
+			assert col in Output.table.keys()
+			w += Output.table[col][3]
+
+		if ((self.col) and (w > self.width)):
+			print("Minimum width required to print {} = {}. Exiting.".format(",".join(self.cols), w))
+			sys.exit(1)
+
+		remainder = self.width - w
+
+		if ("kernel" in self.cols) and ("params" in self.cols):
+			Output.table["kernel"][3] = int(remainder/2)
+			Output.table["params"][3] = int(remainder/2)
+		elif ("kernel" in self.cols):
+			Output.table["kernel"][3] = remainder
+		elif ("params" in self.cols):
+			Output.table["params"][3] = remainder
+
+		#header format
+		cadena = ""
+		for col in self.cols:
+			_,_,t,w = Output.table[col]
+			cadena += "%-{}.{}s ".format(w,w)
+
+		self.hFormat = cadena
+
+		#data format
+		cadena = ""
+		for col in self.cols:
+			_,_,t,w = Output.table[col]
+			if (t == str):
+				cadena += "%-{}.{}s ".format(w,w)
+			elif (t == int):
+				cadena += "%{}d ".format(w)
+
+		self.dFormat = cadena
+
+	def foo(self, cadena, pformat):
+		if self.csv:
+			cadena = ",".join(map(lambda x : '"' + str(x) + '"', cadena))
+		elif self.col:
+			cadena = pformat % cadena
+		else:
+			cadena = " ".join(map(str,cadena))
+
+		try:
+			print(cadena)
+		except IOError as e:
+			#gracefully handle pipes
+			if e.errno == errno.EPIPE:
+				# Python flushes standard streams on exit; redirect remaining output
+				# to devnull to avoid another BrokenPipeError at shutdown
+
+				devnull = os.open(os.devnull, os.O_WRONLY)
+				os.dup2(devnull, sys.stdout.fileno())
+				sys.exit(0)
+			else:
+				sys.exit(-1)
+
+	def header(self):
+		cadena = ()
+		for col in self.cols:
+			h = Output.table[col][0]
+			cadena = cadena + (h,)
+
+		self.foo(cadena, self.hFormat)
+
+	def data(self, a):
+		if a.dir == "":
+			direc = "na"
+		else:
+			direc = a.dir
+
+		if a.op == "":
+			op = "na"
+		else:
+			op = a.op
+
+		if a.mod == "":
+			mod = "na"
+		else:
+			mod = a.mod
+
+		cadena = ()
+		for col in self.cols:
+			attr = Output.table[col][1]
+			val = getattr(a, attr)
+
+			if col == "layer":
+				assert(type(val) == list)
+				val = ":".join(val)
+				val = "-" if val == "" else val
+
+			if col == "trace":
+				assert(type(val) == list)
+				if self.col and len(val):
+					val = val[-1]
+					val = val.split("/")[-1]
+				else:
+					val = ",".join(val)
+					val = "-" if val == "" else val
+
+			if col in ["seq", "altseq"]:
+				assert(type(val) == list)
+				val = ",".join(map(str,val))
+				val = "-" if val == "" else val
+
+			cadena = cadena + (val,)
+	
+		self.foo(cadena, self.dFormat)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/pointwise.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/pointwise.py
new file mode 100644
index 0000000000..3c9afc52a7
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/pointwise.py
@@ -0,0 +1,166 @@
+import numpy as np
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class Pointwise(OperatorLayerBase):
+
+	ops = []
+	ops += ["__abs__", "__neg__", "__invert__"]
+	ops += ["__add__", "__sub__", "__mul__", "__floordiv__", "__truediv__", "__pow__", "__mod__"]
+	ops += ["__radd__", "__rsub__", "__rmul__", "__rdiv__", "__rtruediv__", "__rfloordiv__", "__rpow__"]
+	ops += ["__iadd__", "__isub__", "__imul__", "__itruediv__",]
+	ops += ["__lt__", "__gt__", "__ge__", "__le__", "__eq__", "__ne__",]
+	ops += ["lt", "lt_", "gt", "gt_", "ge", "ge_", "le", "le_", "eq", "eq_", "ne", "ne_",]
+	ops += ["__and__", "__or__", "__xor__", "__lshift__", "__rshift__"]
+	ops += ["__iand__", "__ior__", "__ixor__", "__ilshift__", "__irshift__"]
+	ops += ["abs", "abs_", "neg", "neg_"]
+	ops += ["add", "add_", "div", "div_", "mul", "mul_", "reciprocal", "reciprocal_", "remainder", "remainder_", "sub", "sub_",]
+	ops += ["addcdiv", "addcdiv_", "addcmul", "addcmul_"]
+	ops += ["exp", "exp_", "exp1m", "exp1m_", "log", "log_", "log10", "log10_", "log1p", "log1p_", "log2", "log2_", "pow", "pow_", "rsqrt", "rsqrt_", "sqrt", "sqrt_",]
+	ops += ["ceil", "ceil_", "clamp", "clamp_", "floor", "floor_", "fmod", "fmod_", "frac", "frac_", "round", "round_", "sign", "sign_", "trunc", "trunc_"]
+	ops += ["acos", "acos_", "asin", "asin_", "atan", "atan_", "atan2", "atan2_", "cos", "cos_", "cosh", "cosh_", "sin", "sin_", "sinh", "sinh_", "tan", "tan_", "sigmoid", "sigmoid_", "tanh", "tanh_"]
+	ops += ["digamma", "erf", "erf_", "erfc", "erfc_", "erfinv", "erfinv_", "lerp", "lerp_", "mvlgamma",]
+
+	@staticmethod
+	def foo(d):
+		return d['name'],d['type'],d['shape'],d['dtype']
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		self.dir = d.dir
+		assert (d.dir in ["fprop", "bprop"])
+		assert (op in Pointwise.ops)
+
+		#Filter out all named parameters (kwargs).
+		#This might require revisiting in future.
+		args = list(filter(lambda x : x['name'] == "", args))
+
+		#Filter out non tensors
+		args = list(filter(lambda x : x['type'] == "tensor", args))
+
+		if (len(args) == 0):
+			self.shape = [(1,)]
+			self.type = "float32" #FIX
+
+		elif (len(args) == 1):
+			in0 = args[0]
+			_,t0,s0,dt0 = Pointwise.foo(in0)
+			assert (t0 == "tensor")
+			self.shape = [s0,]
+			self.type = dt0
+
+		elif (len(args) == 2):
+			in0,in1 = args
+			_,t0,s0,dt0 = Pointwise.foo(in0)
+			_,t1,s1,dt1 = Pointwise.foo(in1)
+			assert (t0 == t1 == "tensor")
+			assert (dt0 == dt1)
+			self.shape = [s0,s1]
+			self.type = dt0
+
+		elif (len(args) == 3):
+			in0,in1,in2 = args
+			_,t0,s0,dt0 = Pointwise.foo(in0)
+			_,t1,s1,dt1 = Pointwise.foo(in1)
+			_,t2,s2,dt2 = Pointwise.foo(in2)
+			assert (t0 == t1 == t2 == "tensor")
+			assert (dt0 == dt1 == dt2)
+			self.shape = [s0,s1,s2]
+			self.type = dt0
+		else:
+			assert False
+		return
+
+	def params(self):
+		p = OrderedDict([('T',self.shape), ('type', self.type)])
+		return p
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def elems(self):
+		tensor = self.shape
+		t = self.type
+
+		if (len(tensor) == 1):
+			elems = 2 * Utility.numElems(tensor[0])
+		elif (len(tensor) == 2):
+			if (tensor[0] == tensor[1]):	# same shape
+				elems = Utility.numElems(tensor[0])
+				if self.dir == "fprop":
+					elems *= 3
+				else:
+					if (self.op_ in ["add", "__add__", "sub", "__sub__", "__isub__"]):
+						elems *= 2
+					elif (self.op_ in ["__mul__", "__rmul__", "div", "__truediv__"]):
+						elems *= 3
+					else:
+						assert False
+			else:	#check for broadcast conditions
+				array1 = np.empty(list(tensor[0]))
+				array2 = np.empty(list(tensor[1]))
+				try:
+					out = np.broadcast(array1, array2).shape
+				except:
+					assert False
+
+				elems = Utility.numElems(tensor[0])
+				elems += Utility.numElems(tensor[1])
+				elems += Utility.numElems(out)
+				#TODO bprop
+		elif (len(tensor) == 3):
+			if (tensor[0] == tensor[1] == tensor[2]):	#same shape
+				elems = Utility.numElems(tensor[0])
+				elems *= 4
+			else:
+				assert False
+		else:
+			assert False
+
+		return elems
+
+	def bytes(self):
+		return self.elems() * Utility.typeToBytes(self.type)
+
+	def flops(self):
+		# Note: some cases may still be missing.
+
+		f = 0
+		if self.op_ in ["__abs__", "__neg__", "__add__", "__sub__", "__mul__",
+					"__radd__", "__rmul__", "__iadd__", "__isub__", "__imul__", "__itruediv__",
+					"abs", "abs_", "neg", "neg_", "add", "add_", "div", "div_", "mul", "mul_",
+					"sub", "sub_", "exp", "exp_", "sign", "sign_", "trunc", "trunc_",
+					"sin", "sin_", "cos", "cos_", "sinh", "sinh_", "cosh", "cosh_",
+					"sqrt", "sqrt_", "rsqrt", "rsqrt_", "__lt__", "__gt__", "__ge__", "__le__",
+					"__eq__", "__ne__", "lt", "lt_", "gt", "gt_", "ge", "ge_", "le", "le_",
+					"eq", "eq_", "ne", "ne_", "ceil", "ceil_", "clamp", "clamp_", "floor", "floor_",
+					"round", "sign", "sign_", "trunc", "trunc_"]:
+			# We're counting only one operand, not two (2 operands, 1 op)
+			f = self.elems() / 2
+		elif self.op_ in ["fmod", "fmod_"]:
+			f = self.elems()
+		elif self.op_ in ["tanh", "tanh_", "sigmoid", "sigmoid_", "log", "log_", "log2",
+			 "log2_", "log10", "log10_"]:
+			f = self.elems() * 2
+		elif self.op_ in ["asin", "asin_", "acos", "acos_", "atan", "atan_"]:
+			# no intrinsic, hence slow execution
+			# surprisingly, asin/acos and atan were all the same (via nvprof measurement)
+			f = self.elems() * 10
+
+		return f
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/pooling.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/pooling.py
new file mode 100644
index 0000000000..3f342b4d4d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/pooling.py
@@ -0,0 +1,59 @@
+from .collections import OrderedDict
+from .utility import Utility
+
+# Work in progress.
+
+#poolFuncs = ["max_pool2d_with_indices_forward", "max_pool2d_with_indices"]
+class MaxPool2d(object):
+
+	def parse(marker):
+
+		def convert2Tuple(arg):
+			assert (arg['type'] in ["int", "tuple"])
+			if arg['type'] == "int":
+				return (arg['value'], arg['value'])
+			else:
+				return arg['value']
+
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+		assert (mod == "torch.nn.functional")
+		assert (op == "max_pool2d")
+		assert (len(args) >= 2)
+
+		#input
+		assert (args[0]['name'] == "")
+		inp = args[0]
+		assert (inp['type'] == "tensor")
+		i = inp['shape']
+		t = inp['dtype']
+		assert (len(i) == 4) #nchw tensor
+
+		#kernel
+		if (args[1]['name'] == ""):
+			k = args[1]
+		else:
+			k = list(filter(lambda x : x['name'] == "kernel_size", args))[0]
+		k = convert2Tuple(k)
+
+		#stride
+		s = k #default value
+		if ((len(args) >= 3) and args[2] == ""):
+			s = args[2]
+			s = convert2Tuple(s)
+		elif any(x['name'] == "stride" for x in args):
+			s = list(filter(lambda x : x['name'] == "stride", args))[0]
+			s = convert2Tuple(s)
+
+		#padding
+		p = (0,0)
+		if ((len(args) >= 4) and args[3] == ""):
+			p = args[3]
+			p = convert2Tuple(p)
+		elif any(x['name'] == "padding" for x in args):
+			p = list(filter(lambda x : x['name'] == "padding", args))[0]
+			p = convert2Tuple(p)
+		
+		params = OrderedDict([('T', i), ('K', k), ('s',s), ('p',p), ('type', t)])
+		return params
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/prof.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/prof.py
new file mode 100644
index 0000000000..a3467e6f3f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/prof.py
@@ -0,0 +1,256 @@
+#!/usr/bin/env python3
+
+"""
+This script reads the output (Python dictionary) created by parse.py.
+For every kernel (line) in the input it determines
+	module / class name e.g. torch.nn.functional
+	operator name e.g. linear
+	kernel parameters e.g. GEMM M, N, K, datatype
+	bytes
+	flops
+	tensor core usage
+	direction (fprop, bprop)
+	and other things. Please see the tool usage.
+"""
+
+from .usage import parseArgs
+from .output import Output
+from .utility import Utility
+from .pointwise import Pointwise
+from .convert import Convert
+from .blas import *
+from .embedding import Embedding
+from .reduction import *
+from .dropout import Dropout
+from .softmax import *
+#from pooling import * # work in progress
+from .linear import Linear
+from .optim import Adam
+from .misc import *
+from .conv import Conv
+from .activation import Activation
+from .index_slice_join_mutate import Cat, Reshape, MaskedScatter, Gather, Nonzero, IndexSelect, MaskedSelect
+from .recurrentCell import RNNCell
+from .normalization import BatchNorm
+from .randomSample import RandPerm
+from .loss import MSELoss
+from .data import Data
+
+def findFpropKernel(seq):
+	#Find the last fprop kernel with the same seqId
+	#First look at seqId and then at altSeqId
+	for idx in reversed(range(len(kernels))):
+		k = kernels[idx]
+		if (seq in k['seqId']) and (k['dir'] == "fprop"):
+			return idx
+
+	for idx in reversed(range(len(kernels))):
+		k = kernels[idx]
+		if (seq in k['altSeqId']) and (k['dir'] == "fprop"):
+			return idx
+
+	return -1
+	#print("Error: seqId {} not found.".format(seq), file=sys.stderr)
+	#assert False
+
+def foo(mod, op, d):
+	if (op[0] == "linear"):
+		xx = Linear(d)
+
+	# rnncell, lstmcell, grucell
+	elif (mod[0] in["LSTMCell", "GRUCell"]) and (op[0] == "forward"):
+		xx = RNNCell(d)
+
+	elif op[0] in ["conv1d", "conv2d",]:
+		xx = Conv(d)
+
+	elif (op[0] in Pointwise.ops):
+		xx = Pointwise(d)
+
+	elif (op[0] in Convert.ops):
+		xx = Convert(d)
+
+	elif op[0] in ["__matmul__", "matmul"]:
+		xx = Matmul(d)
+
+	elif op[0] == "embedding":
+		xx = Embedding(d)
+
+	#reduction
+	elif op[0] == "sum":
+		xx = Sum(d)
+
+	elif op[0] == "mean":
+		xx = Mean(d)
+
+	elif op[0] == "norm":
+		xx = Norm(d)
+
+	elif op[0] == "dropout":
+		xx = Dropout(d)
+
+	#Index, Slice, Join, Mutate
+	elif (op[0] == "cat"):
+		xx = Cat(d)
+
+	elif (op[0] == "reshape"):
+		xx = Reshape(d)
+
+	elif (op[0] == "masked_scatter_"):
+		xx = MaskedScatter(d)
+
+	elif (op[0] == "gather"):
+		xx = Gather(d)
+
+	elif (op[0] == "nonzero"):
+		xx = Nonzero(d)
+
+	elif (op[0] == "index_select"):
+		xx = IndexSelect(d)
+
+	elif (op[0] == "masked_select"):
+		xx = MaskedSelect(d)
+
+	#blas
+	elif op[0] in ["addmm", "addmm_"]:
+		xx = Addmm(d)
+
+	elif op[0] == "mm":
+		xx = Mm(d)
+
+	elif op[0] == "bmm":
+		xx = Bmm(d)
+
+	#softmax
+	elif op[0] == "softmax":
+		xx = Softmax(d)
+
+	elif op[0] == "log_softmax":
+		xx = LogSoftmax(d)
+
+	#loss
+	elif op[0] == "mse_loss":
+		xx = MSELoss(d)
+
+	#optimizers
+	elif op[0] == "adam":
+		xx = Adam(d)
+
+	#normalization
+	elif op[0] == "batch_norm":
+		xx = BatchNorm(d)
+
+	#random
+	elif op[0] == "randperm":
+		xx = RandPerm(d)
+
+	#misc
+	elif op[0] == "copy_":
+		xx = Copy(d)
+
+	elif op[0] == "clone":
+		xx = Clone(d)
+
+	elif op[0] == "contiguous":
+		xx = Contiguous(d)
+
+	elif op[0] == "any":
+		xx = Any(d)
+
+	elif (op[0] in Activation.ops):
+		xx = Activation(d)
+
+	elif op[0] == "to":
+		xx = Convert(d)
+
+	else:
+		xx = Foo(d)
+
+	return xx
+
+def main():
+	#Read cmd line arguments
+	cmdArgs = parseArgs()
+
+	output = Output(cmdArgs)
+	output.header()
+
+	idx = -1
+	#Read in all the kernel info
+	for line in cmdArgs.file:
+		idx += 1
+		kernel = eval(line)
+		assert(kernel)
+		kernels.append(kernel)
+
+		k = kernel
+		d = Data(k)
+
+		mod = k['mod']
+		op = k['op']
+
+		flops = 0
+		params = {"na":"na"}
+		tc = "na"
+		bytes = 0
+
+		if (d.dir == "bprop"):
+			d.seqMarker = k['seqMarker']
+			seq = k['seqId']
+			if len(seq) > 1:
+				pass
+			seq = k['seqId'][:1]
+			assert (len(seq) == 1), seq
+			#assert (seq[0] != 0)
+			assert (len(d.seqMarker) > 0)
+			#If there is no useful marker associated, use the
+			#sequence number to find the kernel from fprop
+			if len(d.argMarker) == 0:
+				index = findFpropKernel(seq[0])
+				if index >= 0:
+					d.argMarker = kernels[index]['marker']
+					d.modMarker = kernels[index]['reprMarkers']
+					mod = kernels[index]['mod']
+					op = kernels[index]['op']
+
+					d.layer = kernels[index]['layer']
+					d.trace = kernels[index]['trace']
+
+		# Check if marker has our annotations
+		if len(d.argMarker) and Utility.hasNVTX(d.argMarker[0]):
+
+			xx = foo(mod, op, d)
+
+			bytes = xx.bytes()
+			flops = xx.flops()
+			op = xx.op()
+			params = xx.params()
+			tc = xx.tc()
+
+		if type(op) is list:
+			if len(op):
+				op = op[0]
+			else:
+				op = ""
+
+		if type(mod) is list:
+			if len(mod):
+				mod = mod[0]
+			else:
+				mod = ""
+
+		d.index = idx+1
+
+		# The following 8 come from operator class functions.
+		d.setParams(params)
+		d.tc = tc
+		d.flops = flops
+		d.bytes = bytes
+		d.mod = mod
+		d.op = op
+
+		output.data(d)
+
+kernels = []
+if __name__ == '__main__':
+	main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/randomSample.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/randomSample.py
new file mode 100644
index 0000000000..f7521bf348
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/randomSample.py
@@ -0,0 +1,43 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class RandPerm(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "torch")
+		assert (op == "randperm")
+		assert (len(args) == 1)
+		n = args[0]
+		assert n['type'] == "int"
+		self.n = n['value']
+
+	def params(self):
+		p = OrderedDict([('N', self.n)])
+		return p
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def bytes(self):
+		return self.n * Utility.typeToBytes("int64")
+
+	def flops(self):
+		# Depends on RNG but this is probably a reasonable assumption.
+		return self.n * 3
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/recurrentCell.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/recurrentCell.py
new file mode 100644
index 0000000000..945a7b158c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/recurrentCell.py
@@ -0,0 +1,207 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+def hasTileSize(name):
+	if ("sgemm" in name) or ("884gemm" in name) or ("hgemm" in name):
+		return True
+	else:
+		return False
+
+def ctaTile(name):
+	name = name.split("_")
+	name = list(filter(lambda x : "x" in x, name))
+	name = list(filter(lambda x : "slice" not in x, name))
+	assert(len(name) == 1)
+	name = name[0].split("x")
+	assert(len(name) == 2)
+	name = list(map(int, name))
+	return name[0], name[1]
+
+class RNNCell(OperatorLayerBase):
+	"""
+	This class supports RNNCell, LSTMCell and GRUCell.
+	"""
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		self.name = d.name
+		self.dir = d.dir
+		self.sub = d.sub
+		self.grid = d.grid
+
+		assert (op == "forward")
+		assert (mod in ["LSTMCell", "GRUCell", "RNNCell"])
+		assert (len(args) in [2,3])
+
+		x,h = args[0],args[1]
+		b1,ii = x['shape']
+		b2,hh = h['shape']
+		assert b1 == b2
+		assert x['dtype'] == h['dtype']
+		t = x['dtype']
+
+		self.cell = mod
+		self.inp = ii
+		self.hid = hh
+		self.b = b1
+		self.type = t
+
+		self.multiple = 1
+		if self.cell == "LSTMCell":
+			self.multiple = 4
+		elif self.cell == "GRUCell":
+			self.multiple = 3
+
+		self.gemm = None
+		self.m = None
+		self.n = None
+		self.k = None
+		self.elems = 0
+
+		self.bar()
+		
+	def params(self):
+		if self.gemm is None:
+			p = OrderedDict([('cell', self.cell), ('X', self.inp), ('H', self.hid), ('B', self.b), ('type', self.type)])
+		else:
+			assert self.m is not None
+			assert self.n is not None
+			assert self.k is not None
+			p = OrderedDict([('gemm', self.gemm), ('M', self.m), ('N', self.n), ('K', self.k), ('type', self.type)])
+		return p
+
+	def tc(self):
+		if "gemm" in self.name:
+			return 1 if "884gemm" in self.name else 0
+		else:
+			return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def bytes(self):
+		if self.gemm is not None:
+			m, n, k, t = self.m, self.n, self.k, self.type
+			b = (m*k + k*n + m*n) * Utility.typeToBytes(t)
+		elif self.elems != 0:
+			b = self.elems * Utility.typeToBytes(self.type)
+		else:
+			b = 0
+		return b
+
+	def flops(self):
+		if self.gemm is not None:
+			m, n, k = self.m, self.n, self.k
+			f = 2*m*n*k
+		elif self.elems != 0:
+			f = 0 #TODO
+		else:
+			f = 0
+		return f
+
+	def bar(self):
+		cell = self.cell
+		X = self.inp
+		H = self.hid
+		B = self.b
+		t = self.type
+		subseqId = self.sub
+		direc = self.dir
+		name = self.name
+		grid = self.grid
+		multiple = self.multiple
+
+		if direc == "fprop":
+			subseqId = subseqId % 3
+			if subseqId == 0: #layer gemm
+				self.gemm = "layer"
+				self.m = multiple*H
+				self.n = B
+				self.k = X
+			elif subseqId == 1: #recurrent gemm
+				self.gemm = "recur"
+				self.m = multiple*H
+				self.n = B
+				self.k = H
+			else:
+				layerGemmElems = multiple*H*B
+				recurGemmElems = multiple*H*B
+				cElems = H*B
+				hElems = H*B
+				totElems = layerGemmElems + recurGemmElems + 2*cElems + hElems
+				self.elems = totElems
+
+		else:
+			if ("gemm" in name) and hasTileSize(name):	#gemm
+				#Get cta tile size
+				tileX, tileY = ctaTile(name)
+				#Get grid dimensions
+				grid = grid.split(",")
+				gridX,gridY,gridZ = map(lambda x : int(x), grid)
+
+				gemmM = tileX * gridX
+				gemmN = tileY * gridY
+
+				if name[-3:] == "_nn": # dgrad
+					if (gemmM == H):	# recurrent dgrad
+						#Ideally gemmN = B, but we have a limited set of tile sizes.
+						gemmN = B
+						gemmK = multiple*H
+
+						self.gemm = "recur"
+						self.m = gemmM
+						self.n = gemmN
+						self.k = gemmK
+
+					elif (gemmM == X):	# layer dgrad
+						#assert(gemmN % B == 0)
+						gemmK = multiple*H
+
+						self.gemm = "layer"
+						self.m = gemmM
+						self.n = gemmN
+						self.k = gemmK
+
+					else:
+						pass
+
+				elif name[-3:] == "_nt": #wgrad
+					if (gemmM == H):	#recurrent wgrad
+						assert (gemmN == multiple*H)
+						gemmK = B
+
+						self.gemm = "recur"
+						self.m = gemmM
+						self.n = gemmN
+						self.k = gemmK
+
+					elif (gemmM == X):	#layer wgrad
+						assert (gemmN == multiple*H)
+						gemmK = B
+
+						self.gemm = "layer"
+						self.m = gemmM
+						self.n = gemmN
+						self.k = gemmK
+
+					else:
+						pass
+				else:
+					pass
+			else:
+				pass
+
+		return
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/reduction.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/reduction.py
new file mode 100644
index 0000000000..af27035235
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/reduction.py
@@ -0,0 +1,150 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class Mean(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod in ["torch", "Tensor"])
+		assert (op == "mean")
+
+		#Filter out named parameters
+		args = list(filter(lambda x : x['name'] == '', args))
+
+		assert (len(args) <= 2)
+		i = args[0]
+
+		self.shape = i['shape']
+		self.type = i['dtype']
+		self.dir = d.dir
+		self.sub = d.sub
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def elems(self):
+		return Utility.numElems(self.shape)
+
+	def bytes(self):
+		if self.sub == 0:
+			return self.elems() * Utility.typeToBytes(self.type)
+		else:
+			return 0
+
+	def flops(self):
+		if self.sub == 0:
+			return self.elems() + 1
+		else:
+			return 0
+
+class Sum(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod in ["torch", "Tensor"])
+		assert (op == "sum")
+		assert (len(args) >= 1)
+
+		#Get input
+		if (args[0]['name'] == ""):
+			i = args[0]
+		else:
+			i = list(filter(lambda x : x['name'] == "input", args))[0]
+
+		self.shape = i['shape']
+		self.type = i['dtype']
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def elems(self):
+		return Utility.numElems(self.shape)
+
+	def flops(self):
+		# Note: This is incorrect, need to calculate actual flops (say via nvprof)
+		return self.elems()
+
+	def bytes(self):
+		return self.elems() * Utility.typeToBytes(self.type)
+
+class Norm(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod in ["torch", "Tensor"])
+		assert (op == "norm")
+		#assert (len(args) == 1)
+		i = args[0]
+		self.shape = i['shape']
+		self.type = i['dtype']
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def elems(self):
+		return Utility.numElems(self.shape)
+
+	def bytes(self):
+		return self.elems() * Utility.typeToBytes(self.type)
+
+	def flops(self):
+		# square and add plus sqrt
+		return 2 * self.elems() + 1
+
+	def tc(self):
+		return "-"
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/softmax.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/softmax.py
new file mode 100644
index 0000000000..4271a8d949
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/softmax.py
@@ -0,0 +1,115 @@
+from collections import OrderedDict
+from .utility import Utility
+from .base import OperatorLayerBase
+
+class Softmax(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "torch.nn.functional")
+		assert (op == "softmax")
+
+		#Filter out named parameters
+		args = list(filter(lambda x : x['name'] == '', args))
+
+		assert (len(args) <= 2)
+		self.shape = args[0]['shape']
+		self.type = args[0]['dtype']
+		self.dir = d.dir
+
+		return
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def tc(self):
+		return "-"
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def elems(self):
+		return Utility.numElems(self.shape)
+
+	def flops(self):
+		# Note: exp, sum-reduce, divide
+		#flops = elems * 3
+		return 0
+
+	def bytes(self):
+		b = self.elems() * Utility.typeToBytes(self.type)
+		b *= 3 if self.dir == "fprop" else 5 #verify
+		return b
+
+class LogSoftmax(OperatorLayerBase):
+
+	def __init__(self, d):
+		marker = eval(d.argMarker[0])
+		mod = marker['mod']
+		op = marker['op']
+		args = marker['args']
+
+		self.marker = marker
+		self.mod_ = mod
+		self.op_ = op
+		self.args = args
+
+		assert (mod == "torch.nn.functional")
+		assert (op == "log_softmax")
+
+		#Filter out named parameters
+		args = list(filter(lambda x : x['name'] == '', args))
+
+		assert (len(args) <= 2)
+
+		#Get input
+		if (args[0]['name'] == ""):
+			i = args[0]
+		else:
+			i = list(filter(lambda x : x['name'] == "input", args))[0]
+
+		t = i['dtype']
+
+		self.shape = i['shape']
+		self.type = i['dtype']
+		self.dir = d.dir
+		return
+
+	def op(self):
+		return self.op_
+
+	def mod(self):
+		return self.mod_
+
+	def tc(self):
+		return "-"
+
+	def params(self):
+		p = OrderedDict([('T', self.shape), ('type', self.type)])
+		return p
+
+	def elems(self):
+		return Utility.numElems(self.shape)
+
+	def flops(self):
+		# Note: exp, sum-reduce, divide, log
+		#flops = elems * 4
+		return 0
+
+	def bytes(self):
+		b = self.elems() * Utility.typeToBytes(self.type)
+		b *= 3 if self.dir == "fprop" else 5 #verify
+		return b
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/usage.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/usage.py
new file mode 100644
index 0000000000..3a299d5650
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/usage.py
@@ -0,0 +1,73 @@
+import sys
+import argparse
+
+def parseArgs():
+	"""
+	Print usage and parse arguments.
+	"""
+
+	def check_cols(value):
+		valid = ["idx", "seq", "altseq", "tid", "layer", "trace", "dir", "sub", "mod", "op", "kernel", "params", "sil", "tc", "device", "stream", "grid", "block", "flops", "bytes"]
+		cols = value.split(",")
+		for col in cols:
+			if col not in valid:
+				raise argparse.ArgumentTypeError("{} is not a valid column name. Valid column names are {}.".format(col, ",".join(valid)))
+		return cols
+
+	def openFile(f):
+		try:
+			d = open(f, "r")
+			return d
+		except IOError:
+			print("Error opening file {}. Exiting.".format(f), file=sys.stderr)
+			sys.exit(1)
+
+	parser = argparse.ArgumentParser(prog=sys.argv[0], description="PyTorch Profiler", formatter_class=argparse.RawTextHelpFormatter)
+	parser.add_argument("file",
+		nargs='?',
+		type=str,
+		default=None,
+		help="Output of parse.py (Python dictionary).")
+
+	parser.add_argument("-c",
+		type=check_cols,
+		default="idx,dir,sub,mod,op,kernel,params,sil",
+		help='''Comma seperated names of columns to print.
+idx:      Index
+seq:      PyTorch Sequence Id
+altseq:   PyTorch Alternate Sequence Id
+tid:      Thread Id
+layer:    User annotated NVTX string (can be nested)
+trace:    Function Call Trace
+dir:      Direction
+sub:      Sub Sequence Id
+mod:      Module
+op:       Operattion
+kernel:   Kernel Name
+params:   Parameters
+sil:      Silicon Time (in ns)
+tc:       Tensor Core Usage
+device:   GPU Device Id
+stream:   Stream Id
+grid:     Grid Dimensions
+block:    Block Dimensions
+flops:    Floating point ops (FMA = 2 FLOPs)
+bytes:    Number of bytes in and out of DRAM
+e.g. -c idx,kernel,sil''')
+
+	group = parser.add_mutually_exclusive_group()
+	group.add_argument("--csv",
+		action="store_true",
+		default=False,
+		help="Print a CSV output.")
+	group.add_argument("-w",
+		type=int,
+		default=0,
+		help="Width of columnated output.")
+
+	args = parser.parse_args()
+	if args.file is None:
+		args.file = sys.stdin
+	else:
+		args.file = openFile(args.file)
+	return args
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/utility.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/utility.py
new file mode 100644
index 0000000000..d2bede357b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/pyprof/prof/utility.py
@@ -0,0 +1,60 @@
+from functools import reduce
+
+class Utility(object):
+
+	@staticmethod
+	def numElems(shape):
+		assert (type(shape) == tuple)
+		return reduce(lambda x,y: x*y, shape, 1)
+
+	@staticmethod
+	def typeToBytes(t):
+		if (t in ["uint8", "int8", "byte", "char", "bool"]):
+			return 1
+		elif (t in ["float16", "half", "int16", "short"]):
+			return 2
+		elif (t in ["float32", "float", "int32", "int"]):
+			return 4
+		elif (t in ["int64", "long", "float64", "double"]):
+			return 8
+		assert False
+
+	@staticmethod
+	def typeToString(t):
+		if (t in ["uint8", "byte", "char",]):
+			return "uint8"
+		elif (t in ["int8",]):
+			return "int8"
+		elif (t in ["int16", "short",]):
+			return "int16"
+		elif (t in ["float16", "half"]):
+			return "fp16"
+		elif (t in ["float32", "float"]):
+			return "fp32"
+		elif (t in ["int32", "int",]):
+			return "int32"
+		elif (t in ["int64", "long"]):
+			return "int64"
+		elif (t in ["float64", "double",]):
+			return "fp64"
+		elif (t in ["bool",]):
+			return "bool"
+		assert False
+
+	@staticmethod
+	def hasNVTX(marker):
+		if type(marker) is str:
+			try:
+				marker = eval(marker)
+			except:
+				return False
+
+		if type(marker) is dict:
+			keys  = marker.keys()
+			return ("mod" in keys) and ("op" in keys) and ("args" in keys)
+		else:
+			return False
+
+	@staticmethod
+	def isscalar(t):
+		return (t in ["float", "int"])
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/README.md
new file mode 100644
index 0000000000..9e86fd8fc1
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/README.md
@@ -0,0 +1 @@
+Under construction...
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/__init__.py
new file mode 100644
index 0000000000..482e01ae21
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/__init__.py
@@ -0,0 +1,127 @@
+from .weight_norm import WeightNorm
+from .reparameterization import Reparameterization
+
+def apply_weight_norm(module, name='', dim=0, hook_child=True):
+    r"""
+    Applies weight normalization to a parameter in the given module.
+    If no parameter is provided, applies weight normalization to all
+    parameters in model (except 1-d vectors and scalars).
+
+    .. math::
+         \mathbf{w} = g \dfrac{\mathbf{v}}{\|\mathbf{v}\|}
+
+    Weight normalization is a reparameterization that decouples the magnitude
+    of a weight tensor from its direction. This replaces the parameter specified
+    by `name` (e.g. "weight") with two parameters: one specifying the magnitude
+    (e.g. "weight_g") and one specifying the direction (e.g. "weight_v").
+    Weight normalization is implemented via a hook that recomputes the weight
+    tensor from the magnitude and direction before every :meth:`~Module.forward`
+    call.
+
+    By default, with `dim=0`, the norm is computed independently per output
+    channel/plane. To compute a norm over the entire weight tensor, use
+    `dim=None`.
+
+    See https://arxiv.org/abs/1602.07868
+
+    Args:
+        module (nn.Module): containing module
+        name (str, optional): name of weight parameter
+        dim (int, optional): dimension over which to compute the norm
+        hook_child (boolean, optional): adds reparameterization hook to direct parent of the 
+            parameters. If False, it's added to `module` instead. Default: True
+
+    Returns:
+        The original module with the weight norm hook
+
+    Example::
+
+        >>> m = apply_weight_norm(nn.Linear(20, 40), name='weight')
+        Linear (20 -> 40)
+        >>> m.weight_g.size()
+        torch.Size([40, 1])
+        >>> m.weight_v.size()
+        torch.Size([40, 20])
+
+    """
+    return apply_reparameterization(module, reparameterization=WeightNorm, hook_child=hook_child,
+                                    name=name, dim=dim)
+
+def remove_weight_norm(module, name='', remove_all=False):
+    """
+    Removes the weight normalization reparameterization of a parameter from a module.
+    If no parameter is supplied then all weight norm parameterizations are removed.
+    Args:
+        module (nn.Module): containing module
+        name (str, optional): name of weight parameter
+    Example:
+        >>> m = apply_weight_norm(nn.Linear(20, 40))
+        >>> remove_weight_norm(m)
+    """
+    return remove_reparameterization(module, reparameterization=WeightNorm,
+                                    name=name, remove_all=remove_all)
+
+def apply_reparameterization(module, reparameterization=None, name='', dim=0, hook_child=True):
+    """
+    Applies a given weight reparameterization (such as weight normalization) to
+    a parameter in the given module. If no parameter is given, applies the reparameterization
+    to all parameters in model (except 1-d vectors and scalars).
+
+    Args:
+        module (nn.Module): containing module
+        reparameterization (Reparameterization): reparamaterization class to apply
+        name (str, optional): name of weight parameter
+        dim (int, optional): dimension over which to perform reparameterization op
+        hook_child (boolean, optional): adds reparameterization hook to direct parent of the 
+            parameters. If False, it's added to `module` instead. Default: True
+
+    Returns:
+        The original module with the reparameterization hook
+
+    Example::
+
+        >>> m = apply_reparameterization(nn.Linear(20, 40), WeightNorm)
+        Linear (20 -> 40)
+
+    """
+    assert reparameterization is not None
+    if name != '':
+        Reparameterization.apply(module, name, dim, reparameterization, hook_child)
+    else:
+        names = list(module.state_dict().keys())
+        for name in names:
+            apply_reparameterization(module, reparameterization, name, dim, hook_child)
+    return module
+
+def remove_reparameterization(module, reparameterization=Reparameterization,
+                                name='', remove_all=False):
+    """
+    Removes the given reparameterization of a parameter from a module.
+    If no parameter is supplied then all reparameterizations are removed.
+    Args:
+        module (nn.Module): containing module
+        reparameterization (Reparameterization): reparamaterization class to apply
+        name (str, optional): name of weight parameter
+        remove_all (bool, optional): if True, remove all reparamaterizations of given type. Default: False
+    Example:
+        >>> m = apply_reparameterization(nn.Linear(20, 40),WeightNorm)
+        >>> remove_reparameterization(m)
+    """
+    if name != '' or remove_all:
+        to_remove = []
+        for k, hook in module._forward_pre_hooks.items():
+            if isinstance(hook, reparameterization) and (hook.name == name or remove_all):
+                hook.remove(module)
+                to_remove.append(k)
+        if len(to_remove) > 0:
+            for k in to_remove:
+                del module._forward_pre_hooks[k]
+            return module
+        if not remove_all:
+            raise ValueError("reparameterization of '{}' not found in {}"
+                             .format(name, module))
+    else:
+        modules = [module]+[x for x in module.modules()]
+        for m in modules:
+            remove_reparameterization(m, reparameterization=reparameterization, remove_all=True)
+        return module
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/reparameterization.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/reparameterization.py
new file mode 100644
index 0000000000..24b11ba392
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/reparameterization.py
@@ -0,0 +1,151 @@
+import torch
+from torch.nn.parameter import Parameter
+import sys
+class Reparameterization(object):
+    """
+    Class interface for performing weight reparameterizations
+    Arguments:
+        name (str): name of weight parameter
+        dim (int): dimension over which to compute the norm
+        module (nn.Module): parent module to which param `name` is registered to
+        retain_forward (bool, optional): if False deletes weight on call to 
+            module.backward. Used to avoid memory leaks with DataParallel Default: True
+    Attributes:
+        reparameterization_names (list, str): contains names of all parameters 
+            needed to compute reparameterization.
+        backward_hook_key (int): torch.utils.hooks.RemovableHandle.id for hook used in module backward pass.
+    """
+
+    def __init__(self, name, dim, module, retain_forward=True):
+        self.name = name
+        self.dim = dim
+        self.evaluated = False
+        self.retain_forward = retain_forward
+        self.reparameterization_names = []
+        self.backward_hook_key = None
+        self.module = module
+
+    def compute_weight(self, module=None, name=None):
+        """
+        Computes reparameterized weight value to assign value to module attribute
+        with name `name`.
+        See WeightNorm class for example.
+        Arguments:
+            module (nn.Module): module with weight we'd like to reparameterize
+        Returns:
+            w (Tensor): Tensor object containing value of reparameterized weight
+        """
+        raise NotImplementedError
+
+    def reparameterize(self, name, weight, dim):
+        """
+        Creates Parameters to be used for reparameterization and creates names that
+        for attributes for the module these Parameters will correspond to.
+        The parameters will be registered according to the names provided.
+        See WeightNorm class for example.
+        Arguments:
+            module (nn.Module): module with weight we'd like to reparameterize
+            name (str, optional): name of weight parameter
+            dim (int, optional): dimension over which to compute parameterization
+        Returns:
+            names (list, str): names of Parameters to be used for reparameterization
+            params (list, Parameter): Parameters to be used for reparameterization
+        """
+        raise NotImplementedError
+
+    @staticmethod
+    def apply(module, name, dim, reparameterization=None, hook_child=True):
+        """
+        Applies reparametrization to module's `name` parameter and modifies instance attributes as appropriate.
+        `hook_child` adds reparameterization hook to direct parent of the parameters. If False, it's added to `module` instead.
+        """
+        if reparameterization is None:
+            reparameterization = Reparameterization
+        module2use, name2use = Reparameterization.get_module_and_name(module, name)
+        # does not work on sparse
+        if name2use is None or isinstance(module2use, (torch.nn.Embedding, torch.nn.EmbeddingBag)):
+            return
+
+        if hook_child:
+            fn = reparameterization(name2use, dim, module2use)
+        else:
+            fn = reparameterization(name, dim, module)
+
+        weight = getattr(module2use, name2use)
+        if weight.dim() <= 1:
+            return
+
+        # remove weight from parameter list
+        del module2use._parameters[name2use]
+
+        # add parameters of reparameterization of parameter to module
+        names, params = fn.reparameterize(name2use, weight, dim)
+        for n, p in zip(names, params):
+            module2use.register_parameter(n, p)
+
+        # add parameters to reparameterization so they can be removed later
+        fn.reparameterization_names = names
+
+        setattr(module2use, name2use, None)
+
+        hook_module = module2use
+        if not hook_child:
+            hook_module = module
+        # recompute weight before every forward()
+        hook_module.register_forward_pre_hook(fn)
+
+        # remove weight during backward
+        handle = hook_module.register_backward_hook(fn.backward_hook)
+        # get hook key so we can delete it later
+        fn.backward_hook_key = handle.id
+
+        return fn
+
+    @staticmethod
+    def get_module_and_name(module, name):
+        """
+        recursively fetches (possible) child module and name of weight to be reparameterized
+        """
+        name2use = None
+        module2use = None
+        names = name.split('.')
+        if len(names) == 1 and names[0] != '':
+            name2use = names[0]
+            module2use = module
+        elif len(names) > 1:
+            module2use = module
+            name2use = names[0]
+            for i in range(len(names)-1):
+                module2use = getattr(module2use, name2use)
+                name2use = names[i+1]
+        return module2use, name2use
+
+    def get_params(self, module):
+        """gets params of reparameterization based on known attribute names"""
+        return [getattr(module, n) for n in self.reparameterization_names]
+
+    def remove(self, module):
+        """removes reparameterization and backward hook (does not remove forward hook)"""
+        module2use, name2use = Reparameterization.get_module_and_name(module, self.name)
+        for p in self.get_params(module2use):
+            p.requires_grad = False
+        weight = self.compute_weight(module2use, name2use)
+        delattr(module2use, name2use)
+        for n in self.reparameterization_names:
+            del module2use._parameters[n]
+        module2use.register_parameter(name2use, Parameter(weight.data))
+        del module._backward_hooks[self.backward_hook_key]
+
+    def __call__(self, module, inputs):
+        """callable hook for forward pass"""
+        module2use, name2use = Reparameterization.get_module_and_name(module, self.name)
+        _w = getattr(module2use, name2use)
+        if not self.evaluated or _w is None:
+            setattr(module2use, name2use, self.compute_weight(module2use, name2use))
+            self.evaluated = True
+
+    def backward_hook(self, module, grad_input, grad_output):
+        """callable hook for backward pass"""
+        module2use, name2use = Reparameterization.get_module_and_name(module, self.name)
+        wn = getattr(module2use, name2use)
+        self.evaluated = False
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/weight_norm.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/weight_norm.py
new file mode 100644
index 0000000000..ba54ecc39b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/reparameterization/weight_norm.py
@@ -0,0 +1,78 @@
+import torch
+from torch.nn.parameter import Parameter
+from ..fp16_utils import Fused_Weight_Norm
+import time
+
+from .reparameterization import Reparameterization
+
+def _norm(p, dim):
+    """Computes the norm over all dimensions except dim"""
+    if dim is None:
+        return p.norm()
+    elif dim == 0:
+        output_size = (p.size(0),) + (1,) * (p.dim() - 1)
+        return p.contiguous().view(p.size(0), -1).norm(dim=1).view(*output_size)
+    elif dim == p.dim() - 1:
+        output_size = (1,) * (p.dim() - 1) + (p.size(-1),)
+        return p.contiguous().view(-1, p.size(-1)).norm(dim=0).view(*output_size)
+    return _norm(p.transpose(0, dim), 0).transpose(0, dim)
+
+HALF_TYPES = (torch.cuda.HalfTensor, torch.HalfTensor)
+
+class WeightNorm(Reparameterization):
+    r"""
+    Weight normalization is a reparameterization that decouples the magnitude
+    of a weight tensor from its direction. This replaces the parameter specified
+    by `name` (e.g. "weight") with two parameters: one specifying the magnitude
+    (e.g. "weight_g") and one specifying the direction (e.g. "weight_v").
+    Weight normalization is implemented via a hook that recomputes the weight
+    tensor from the magnitude and direction before every :meth:`~Module.forward`
+    call.
+
+    .. math::
+         \mathbf{w} = g \dfrac{\mathbf{v}}{\|\mathbf{v}\|}
+
+    By default, with `dim=0`, the norm is computed independently per output
+    channel/plane. To compute a norm over the entire weight tensor, use
+    `dim=None`.
+    """
+    def compute_weight(self, module=None, name=None):
+        """
+        Computes weight normalized weight value to assign value to module attribute
+        with name `name`.
+        Arguments:
+            module (nn.Module): module with weight we'd like to reparameterize
+        Returns:
+            w (Tensor): Tensor object containing value of reparameterized weight
+        """
+        if module is None:
+            module = self.module
+        if name is None:
+            name = self.name
+        module, name = Reparameterization.get_module_and_name(module, name)
+        g = getattr(module, name + '_g')
+        v = getattr(module, name + '_v')
+
+        fused_weight_norm = Fused_Weight_Norm.apply
+        v = v.contiguous()
+        w = fused_weight_norm(v, g, self.dim)
+
+        return w
+
+    def reparameterize(self, name, weight, dim):
+        """
+        Creates Parameters v and gto be used for weight normalization
+        and creates names that for attributes for the module these Parameters
+        will correspond to. The parameters will be registered according to the names
+        provided.
+        Arguments:
+            module (nn.Module): module with weight we'd like to reparameterize
+            name (str, optional): name of weight parameter
+            dim (int, optional): dimension over which to compute parameterization
+        Returns:
+            names (list, str): names of Parameters to be used for reparameterization
+            params (list, Parameter): Parameters to be used for reparameterization
+        """
+        names = [name + '_g', name + '_v']
+        params = [Parameter(_norm(weight, dim).data), Parameter(weight.data)]
+        return names, params
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/README.md
new file mode 100644
index 0000000000..7383f65eb9
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/README.md
@@ -0,0 +1,81 @@
+# apex.transformer
+
+`apex.transformer` is a module which enables efficient large Transformer models at scale.
+
+`apex.transformer.tensor_parallel` and `apex.transformer.pipeline_parallel` are both based on [NVIDIA/Megatron-LM](https://github.com/NVIDIA/Megatron-LM)'s module.
+The former is based on `megatron.mpu` and the latter is on `megatron.schedules` and `megatron.p2p_communication`.
+
+## Tensor Model Parallel (TP)
+
+APEX's tensor model parallel utilities provides some `torch.nn.Module`'s, custom fused kernels, and PRNG state handling.
+See Appendix B.2 of [Megatron-LM: Training Multi-Billion Parameter Language Models Using Model Parallelism](https://arxiv.org/abs/1909.08053) for the details of
+PRNG state handling.
+
+## Pipeline Model Parallel (PP)
+APEX's pipeline model parallel functions require models to have `.set_input_tensor` because
+the input tensor for `.forward` method can be `None`.
+
+The following is a really casual sketch of training script with apex pp.
+
+```python
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+from apex.transformer import parallel_state
+from apex.transformer.pipeline_parallel import get_forward_backward_func
+
+
+class Model(nn.Module):
+
+    ...
+
+    def __init__(self, *args, **kwargs):
+        super().__init__()
+        pre_process = kwargs.pop("pre_process")
+        post_process = kwargs.pop("post_process")
+
+    def set_input_tensor(self, tensor):
+        self.input_tensor = tensor
+
+    def forward(self, x, ...):
+        if parallel_state.is_pipeline_first_stage():
+            input = x
+        else:
+            input = self.input_tensor
+        ...
+
+
+def model_provider_func(*args, **kwargs):
+    return Model(*args, **kwargs)
+
+
+def loss_func(pred, label):
+    loss = ...
+    averaged_loss = average_losses_across_data_parallel_group([loss])
+    return loss, {'nice_loss': averaged_loss}
+
+
+def forward_step_func(batch, model):
+    input, label = process_batch(batch)
+    out = model(input)
+    return out, partial(loss_func, label)
+
+
+forward_backward_func = get_forward_backward_func(virtual_pipeline_model_parallel_size, pipeline_model_parallel_size)
+
+
+parallel_state.initialize_model_parallel(
+    tensor_model_parallel_size,
+    pipeline_model_parallel_size,
+    virtual_pipeline_model_parallel_size,
+)
+# The following line basically is equivalent to `build_model(Model, wrap_with_ddp, virtual_pipeline_model_parallel_size, *model_args, **model_kwargs)`
+model = build_model(model_provider_func, wrap_with_ddp, virtual_pipeline_model_parallel_size, *model_args, **model_kwargs)
+optimizer = ...
+data_loader = ...
+for epoch in range(num_epochs):
+    for batch in data_loader:
+        forward_backward_func(forward_step_func, batch, model, forward_only=False, tensor_shape)
+        optimizer.step()
+```
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/__init__.py
new file mode 100644
index 0000000000..83528b6622
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/__init__.py
@@ -0,0 +1,21 @@
+from apex.transformer import functional
+from apex.transformer import parallel_state
+from apex.transformer import pipeline_parallel
+from apex.transformer import tensor_parallel
+from apex.transformer import utils
+from apex.transformer.enums import LayerType
+from apex.transformer.enums import AttnType
+from apex.transformer.enums import AttnMaskType
+
+
+__all__ = [
+    "functional",
+    "parallel_state",
+    "pipeline_parallel",
+    "tensor_parallel",
+    "utils",
+    # enums.py
+    "LayerType",
+    "AttnType",
+    "AttnMaskType",
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/enums.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/enums.py
new file mode 100644
index 0000000000..409db846b9
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/enums.py
@@ -0,0 +1,30 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import enum
+
+
+class LayerType(enum.Enum):
+    encoder = 1
+    decoder = 2
+
+
+class AttnType(enum.Enum):
+    self_attn = 1
+    cross_attn = 2
+
+
+class AttnMaskType(enum.Enum):
+    padding = 1
+    causal = 2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/functional/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/functional/__init__.py
new file mode 100644
index 0000000000..d770c88592
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/functional/__init__.py
@@ -0,0 +1,5 @@
+from apex.transformer.functional.fused_softmax import FusedScaleMaskSoftmax
+
+__all__ = [
+    "FusedScaleMaskSoftmax",
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/functional/fused_softmax.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/functional/fused_softmax.py
new file mode 100644
index 0000000000..973725cdee
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/functional/fused_softmax.py
@@ -0,0 +1,199 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import torch
+
+from apex._autocast_utils import _cast_if_autocast_enabled
+from apex.transformer.enums import AttnMaskType
+
+
+class ScaledUpperTriangMaskedSoftmax(torch.autograd.Function):
+    """
+    Fused operation which performs following three operations in sequence
+    1. Scale the tensor.
+    2. Apply upper triangular mask (typically used in gpt models).
+    3. Perform softmax.
+    """
+
+    @staticmethod
+    def forward(ctx, inputs, scale):
+        import scaled_upper_triang_masked_softmax_cuda
+
+        scale_t = torch.tensor([scale])
+        softmax_results = scaled_upper_triang_masked_softmax_cuda.forward(inputs, scale_t[0])
+
+        ctx.save_for_backward(softmax_results, scale_t)
+        return softmax_results
+
+    @staticmethod
+    def backward(ctx, output_grads):
+        import scaled_upper_triang_masked_softmax_cuda
+
+        softmax_results, scale_t = ctx.saved_tensors
+        input_grads = scaled_upper_triang_masked_softmax_cuda.backward(output_grads, softmax_results, scale_t[0])
+
+        return input_grads, None
+
+
+def scaled_upper_triang_masked_softmax(inputs, _, scale):
+    b, np, sq, sk = inputs.size()
+    assert sq == sk, "causal mask is only for self attention"
+    # Reshaping input to 3D tensor (attn_batches, sq, sk)
+    inputs = inputs.view(-1, sq, sk)
+    args = _cast_if_autocast_enabled(inputs, scale)
+    with torch.cuda.amp.autocast(enabled=False):
+        probs = ScaledUpperTriangMaskedSoftmax.apply(*args)
+    return probs.view(b, np, sq, sk)
+
+
+# NOTE (mkozuki): `ScaledMaskedSoftmax` somehow doesn't work well with `torch.cuda.amp.custom_fwd`.
+# Without `cast_inputs` kwarg, somehow inputs are not cast to dtype used in the autocast context.
+# So I needed to manually write two `torch.autograd.Function` inheritances.
+# Fused operation which performs following three operations in sequence
+# 1. Scale the tensor.
+# 2. Apply the mask.
+# 3. Perform softmax.
+class ScaledMaskedSoftmax(torch.autograd.Function):
+
+    @staticmethod
+    def forward(ctx, inputs, mask, scale):
+        import scaled_masked_softmax_cuda
+        scale_t = torch.tensor([scale])
+
+        softmax_results = scaled_masked_softmax_cuda.forward(inputs, mask, scale_t[0])
+        ctx.save_for_backward(softmax_results, scale_t)
+        return softmax_results
+
+    @staticmethod
+    def backward(ctx, output_grads):
+        import scaled_masked_softmax_cuda
+
+        softmax_results, scale_t = ctx.saved_tensors
+
+        input_grads = scaled_masked_softmax_cuda.backward(output_grads, softmax_results, scale_t[0])
+        return input_grads, None, None
+
+
+def scaled_masked_softmax(inputs, mask, scale):
+    # input is 4D tensor (b, np, sq, sk)
+    args = _cast_if_autocast_enabled(inputs, mask, scale)
+    with torch.cuda.amp.autocast(enabled=False):
+        return ScaledMaskedSoftmax.apply(*args)
+
+
+class FusedScaleMaskSoftmax(torch.nn.Module):
+    """
+    fused operation: scaling + mask + softmax
+
+    Arguments:
+        input_in_fp16: flag to indicate if input in fp16 data format.
+        input_in_bf16: flag to indicate if input in bf16 data format.
+        attn_mask_type: attention mask type (pad or causal)
+        scaled_masked_softmax_fusion: flag to indicate user want to use softmax fusion
+        mask_func: mask function to be applied.
+        softmax_in_fp32: if true, softmax in performed at fp32 precision.
+        scale: scaling factor used in input tensor scaling.
+    """
+
+    def __init__(
+        self,
+        input_in_fp16,
+        input_in_bf16,
+        attn_mask_type,
+        scaled_masked_softmax_fusion,
+        mask_func,
+        softmax_in_fp32,
+        scale,
+    ):
+        super().__init__()
+        self.input_in_fp16 = input_in_fp16
+        self.input_in_bf16 = input_in_bf16
+        if self.input_in_fp16 and self.input_in_bf16:
+            raise RuntimeError("both fp16 and bf16 flags cannot be active at the same time.")
+        self.input_in_float16 = self.input_in_fp16 or self.input_in_bf16
+        self.attn_mask_type = attn_mask_type
+        self.scaled_masked_softmax_fusion = scaled_masked_softmax_fusion
+        self.mask_func = mask_func
+        self.softmax_in_fp32 = softmax_in_fp32
+        self.scale = scale
+
+        if not (self.scale is None or softmax_in_fp32):
+            raise RuntimeError("softmax should be in fp32 when scaled")
+
+        if self.scaled_masked_softmax_fusion:
+            if self.attn_mask_type == AttnMaskType.causal:
+                self.fused_softmax_func = scaled_upper_triang_masked_softmax
+            elif self.attn_mask_type == AttnMaskType.padding:
+                self.fused_softmax_func = scaled_masked_softmax
+            else:
+                raise ValueError("Invalid attn_mask_type.")
+
+    def forward(self, input, mask):
+        # [b, np, sq, sk]
+        assert input.dim() == 4
+
+        if self.is_kernel_available(mask, *input.size()):
+            return self.forward_fused_softmax(input, mask)
+        else:
+            return self.forward_torch_softmax(input, mask)
+
+    def is_kernel_available(self, mask, b, np, sq, sk):
+        attn_batches = b * np
+
+        if (
+            self.scaled_masked_softmax_fusion  # user want to fuse
+            and self.input_in_float16  # input must be fp16
+            and mask is not None  # mask tensor must not be None
+            and 16 < sk <= 2048  # sk must be 16 ~ 2048
+            and sq % 4 == 0  # sq must be divisor of 4
+            and attn_batches % 4 == 0  # np * b must be divisor of 4
+        ):
+            if 0 <= sk <= 2048:
+                batch_per_block = self.get_batch_per_block(sq, sk, b, np)
+
+                if self.attn_mask_type == AttnMaskType.causal:
+                    if attn_batches % batch_per_block == 0:
+                        return True
+                else:
+                    if sq % batch_per_block == 0:
+                        return True
+        return False
+
+    def forward_fused_softmax(self, input, mask):
+        # input.shape = [b, np, sq, sk]
+        scale = self.scale if self.scale is not None else 1.0
+        return self.fused_softmax_func(input, mask, scale)
+
+    def forward_torch_softmax(self, input, mask):
+        if self.input_in_float16 and self.softmax_in_fp32:
+            input = input.float()
+
+        if self.scale is not None:
+            input = input * self.scale
+        mask_output = self.mask_func(input, mask) if mask is not None else input
+        probs = torch.nn.Softmax(dim=-1)(mask_output)
+
+        if self.input_in_float16 and self.softmax_in_fp32:
+            if self.input_in_fp16:
+                probs = probs.half()
+            else:
+                probs = probs.bfloat16()
+
+        return probs
+
+    @staticmethod
+    def get_batch_per_block(sq, sk, b, np):
+        import scaled_masked_softmax_cuda
+
+        return scaled_masked_softmax_cuda.get_batch_per_block(sq, sk, b, np)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/microbatches.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/microbatches.py
new file mode 100644
index 0000000000..c97c13c3a8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/microbatches.py
@@ -0,0 +1,172 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+"""Megatron number of micro-batches calculators."""
+from abc import ABC
+from abc import abstractmethod
+from typing import Optional, List
+
+
+def build_num_microbatches_calculator(
+        rank: int,
+        rampup_batch_size: Optional[List[int]],
+        global_batch_size: int,
+        micro_batch_size: int,
+        data_parallel_size: int,
+):
+    # Constant num micro-batches.
+    if rampup_batch_size is None:
+        num_microbatches_calculator = ConstantNumMicroBatches(
+            global_batch_size, micro_batch_size, data_parallel_size
+        )
+        if rank == 0:
+            print(
+                "setting number of micro-batches to constant {}".format(num_microbatches_calculator.get()), flush=True
+            )
+
+    else:
+        assert len(rampup_batch_size) == 3, (
+            "expected the following "
+            "format: --rampup-batch-size <start batch size> "
+            "<batch size incerement> <ramp-up samples>"
+        )
+        start_batch_size = int(rampup_batch_size[0])
+        batch_size_increment = int(rampup_batch_size[1])
+        ramup_samples = int(rampup_batch_size[2])
+        if rank == 0:
+            print(
+                "will use batch size rampup starting from global batch "
+                "size {} to global batch size {} with batch size increments "
+                "{} over {} samples.".format(
+                    start_batch_size, global_batch_size, batch_size_increment, ramup_samples
+                ),
+                flush=True,
+            )
+        num_microbatches_calculator = RampupBatchsizeNumMicroBatches(
+            start_batch_size,
+            batch_size_increment,
+            ramup_samples,
+            global_batch_size,
+            micro_batch_size,
+            data_parallel_size,
+        )
+
+    return num_microbatches_calculator
+
+
+class NumMicroBatchesCalculator(ABC):
+    def __init__(self):
+        self.num_micro_batches = None
+        self.current_global_batch_size = None
+
+    def get(self):
+        return self.num_micro_batches
+
+    def get_current_global_batch_size(self):
+        return self.current_global_batch_size
+
+    @abstractmethod
+    def update(self, consumed_samples, consistency_check):
+        pass
+
+
+class ConstantNumMicroBatches(NumMicroBatchesCalculator):
+    def __init__(self, global_batch_size, micro_batch_size, data_parallel_size):
+        micro_batch_times_data_parallel = micro_batch_size * data_parallel_size
+        assert global_batch_size % micro_batch_times_data_parallel == 0, (
+            "global batch size ({}) is not divisible by micro batch size ({})"
+            " times data parallel size ({})".format(global_batch_size, micro_batch_size, data_parallel_size)
+        )
+        self.num_micro_batches = global_batch_size // micro_batch_times_data_parallel
+        assert self.num_micro_batches >= 1
+        self.current_global_batch_size = global_batch_size
+
+        self.micro_batch_size = micro_batch_size
+
+    def update(self, consumed_samples, consistency_check):
+        pass
+
+
+class RampupBatchsizeNumMicroBatches(NumMicroBatchesCalculator):
+    def __init__(
+        self,
+        start_batch_size,
+        batch_size_increment,
+        ramup_samples,
+        global_batch_size,
+        micro_batch_size,
+        data_parallel_size,
+    ):
+        """Batch size ramp up.
+        Over
+          steps = (global-batch-size - start-batch-size) / batch_size_increment
+        increment batch size from start-batch-size to global-batch-size using
+          rampup-samples / steps
+        samples.
+        Arguments:
+            start_batch_size: global batch size to start with
+            batch_size_increment: global batch size increments
+            ramup_samples: number of samples to use ramp up global
+               batch size from `start_batch_size` to `global_batch_size`
+            global_batch_size: global batch size post rampup
+            micro_batch_size: micro batch size
+            data_parallel_size: data parallel size.
+        """
+
+        self.micro_batch_size = micro_batch_size
+        self.data_parallel_size = data_parallel_size
+        self.micro_batch_times_data_parallel_size = self.micro_batch_size * self.data_parallel_size
+        assert self.micro_batch_times_data_parallel_size > 0
+
+        assert start_batch_size > 0
+        self.start_batch_size = start_batch_size
+
+        assert global_batch_size > 0
+        self.global_batch_size = global_batch_size
+        diff_batch_size = self.global_batch_size - self.start_batch_size
+        assert diff_batch_size >= 0
+        assert batch_size_increment > 0
+        self.batch_size_increment = batch_size_increment
+        assert diff_batch_size % batch_size_increment == 0, (
+            "expected "
+            "global batch size interval ({}) to be divisible by global batch "
+            "size increment ({})".format(diff_batch_size, batch_size_increment)
+        )
+
+        num_increments = diff_batch_size // self.batch_size_increment
+        self.ramup_samples = ramup_samples
+        assert self.ramup_samples >= 0
+        self.rampup_samples_per_increment = self.ramup_samples / num_increments
+
+        # Initialize number of microbatches.
+        self.update(0, False)
+
+    def update(self, consumed_samples, consistency_check):
+
+        if consumed_samples > self.ramup_samples:
+            self.current_global_batch_size = self.global_batch_size
+        else:
+            steps = int(consumed_samples / self.rampup_samples_per_increment)
+            self.current_global_batch_size = self.start_batch_size + steps * self.batch_size_increment
+            assert self.current_global_batch_size <= self.global_batch_size
+
+        if consistency_check:
+            assert self.current_global_batch_size % self.micro_batch_times_data_parallel_size == 0, (
+                "current global "
+                "batch size ({}) is not divisible by micro-batch-size ({}) times"
+                "data parallel size ({})".format(
+                    self.current_global_batch_size, self.micro_batch_size, self.data_parallel_size
+                )
+            )
+        self.num_micro_batches = self.current_global_batch_size // self.micro_batch_times_data_parallel_size
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/parallel_state.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/parallel_state.py
new file mode 100644
index 0000000000..64d4b29190
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/parallel_state.py
@@ -0,0 +1,382 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+"""Model and data parallel groups."""
+import torch
+
+# TODO (mkozuki): Consider dissecting utils as this utils import is here
+# only for ensure_divisibility
+from apex.transformer.utils import ensure_divisibility
+
+
+# Intra-layer model parallel group that the current rank belongs to.
+_TENSOR_MODEL_PARALLEL_GROUP = None
+# Inter-layer model parallel group that the current rank belongs to.
+_PIPELINE_MODEL_PARALLEL_GROUP = None
+# Model parallel group (both intra- and pipeline) that the current rank belongs to.
+_MODEL_PARALLEL_GROUP = None
+# Embedding group.
+_EMBEDDING_GROUP = None
+# Data parallel group that the current rank belongs to.
+_DATA_PARALLEL_GROUP = None
+
+_VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = None
+_VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
+
+# These values enable us to change the mpu sizes on the fly.
+_MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None
+_MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
+_MPU_TENSOR_MODEL_PARALLEL_RANK = None
+_MPU_PIPELINE_MODEL_PARALLEL_RANK = None
+
+# A list of ranks that have a copy of the embedding.
+_EMBEDDING_GLOBAL_RANKS = None
+
+# A list of global ranks for each pipeline group to ease calculation of the source
+# rank when broadcasting from the first or last pipeline stage
+_PIPELINE_GLOBAL_RANKS = None
+
+
+def is_unitialized():
+    """Useful for code segments that may be accessed with or without mpu initialization"""
+    return _DATA_PARALLEL_GROUP is None
+
+
+def initialize_model_parallel(
+    tensor_model_parallel_size_=1, pipeline_model_parallel_size_=1, virtual_pipeline_model_parallel_size_=None
+):
+    """
+    Initialize model data parallel groups.
+
+    Arguments:
+        tensor_model_parallel_size: number of GPUs used to parallelize model tensor.
+        pipeline_model_parallel_size: number of GPUs used to parallelize model pipeline.
+
+    Let's say we have a total of 16 GPUs denoted by g0 ... g15 and we
+    use 2 GPUs to parallelize the model tensor, and 4 GPUs to parallelize
+    the model pipeline. The present function will
+    create 8 tensor model-parallel groups, 4 pipeline model-parallel groups
+    and 8 data-parallel groups as:
+        8 data_parallel groups:
+            [g0, g2], [g1, g3], [g4, g6], [g5, g7], [g8, g10], [g9, g11], [g12, g14], [g13, g15]
+        8 tensor model-parallel groups:
+            [g0, g1], [g2, g3], [g4, g5], [g6, g7], [g8, g9], [g10, g11], [g12, g13], [g14, g15]
+        4 pipeline model-parallel groups:
+            [g0, g4, g8, g12], [g1, g5, g9, g13], [g2, g6, g10, g14], [g3, g7, g11, g15]
+    Note that for efficiency, the caller should make sure adjacent ranks
+    are on the same DGX box. For example if we are using 2 DGX-1 boxes
+    with a total of 16 GPUs, rank 0 to 7 belong to the first box and
+    ranks 8 to 15 belong to the second box.
+    """
+    # Get world size and rank. Ensure some consistencies.
+    assert torch.distributed.is_initialized()
+    world_size = torch.distributed.get_world_size()
+    tensor_model_parallel_size = min(tensor_model_parallel_size_, world_size)
+    pipeline_model_parallel_size = min(pipeline_model_parallel_size_, world_size)
+    ensure_divisibility(world_size, tensor_model_parallel_size * pipeline_model_parallel_size)
+    data_parallel_size = world_size // (tensor_model_parallel_size * pipeline_model_parallel_size)
+    if torch.distributed.get_rank() == 0:
+        print("> initializing tensor model parallel with size {}".format(tensor_model_parallel_size))
+        print("> initializing pipeline model parallel with size {}".format(pipeline_model_parallel_size))
+        print("> initializing data parallel with size {}".format(data_parallel_size))
+
+    num_tensor_model_parallel_groups = world_size // tensor_model_parallel_size
+    num_pipeline_model_parallel_groups = world_size // pipeline_model_parallel_size
+    num_data_parallel_groups = world_size // data_parallel_size
+
+    if virtual_pipeline_model_parallel_size_ is not None:
+        assert pipeline_model_parallel_size_ > 2, \
+            'pipeline-model-parallel size should be greater than 2 with ' \
+            'interleaved schedule'
+        global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
+        global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
+        _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = 0
+        _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = virtual_pipeline_model_parallel_size_
+
+    rank = torch.distributed.get_rank()
+
+    # Build the data-parallel groups.
+    global _DATA_PARALLEL_GROUP
+    assert _DATA_PARALLEL_GROUP is None, "data parallel group is already initialized"
+    all_data_parallel_group_ranks = []
+    for i in range(pipeline_model_parallel_size):
+        start_rank = i * num_pipeline_model_parallel_groups
+        end_rank = (i + 1) * num_pipeline_model_parallel_groups
+        for j in range(tensor_model_parallel_size):
+            ranks = range(start_rank + j, end_rank, tensor_model_parallel_size)
+            all_data_parallel_group_ranks.append(list(ranks))
+            group = torch.distributed.new_group(ranks)
+            if rank in ranks:
+                _DATA_PARALLEL_GROUP = group
+
+    # Build the model-parallel groups.
+    global _MODEL_PARALLEL_GROUP
+    assert _MODEL_PARALLEL_GROUP is None, "model parallel group is already initialized"
+    for i in range(data_parallel_size):
+        ranks = [data_parallel_group_ranks[i] for data_parallel_group_ranks in all_data_parallel_group_ranks]
+        group = torch.distributed.new_group(ranks)
+        if rank in ranks:
+            _MODEL_PARALLEL_GROUP = group
+
+    # Build the tensor model-parallel groups.
+    global _TENSOR_MODEL_PARALLEL_GROUP
+    assert _TENSOR_MODEL_PARALLEL_GROUP is None, "tensor model parallel group is already initialized"
+    for i in range(num_tensor_model_parallel_groups):
+        ranks = range(i * tensor_model_parallel_size, (i + 1) * tensor_model_parallel_size)
+        group = torch.distributed.new_group(ranks)
+        if rank in ranks:
+            _TENSOR_MODEL_PARALLEL_GROUP = group
+
+    # Build the pipeline model-parallel groups and embedding groups
+    # (first and last rank in each pipeline model-parallel group).
+    global _PIPELINE_MODEL_PARALLEL_GROUP
+    global _PIPELINE_GLOBAL_RANKS
+    assert _PIPELINE_MODEL_PARALLEL_GROUP is None, "pipeline model parallel group is already initialized"
+    global _EMBEDDING_GROUP
+    global _EMBEDDING_GLOBAL_RANKS
+    assert _EMBEDDING_GROUP is None, "embedding group is already initialized"
+    for i in range(num_pipeline_model_parallel_groups):
+        ranks = range(i, world_size, num_pipeline_model_parallel_groups)
+        group = torch.distributed.new_group(ranks)
+        if rank in ranks:
+            _PIPELINE_MODEL_PARALLEL_GROUP = group
+            _PIPELINE_GLOBAL_RANKS = ranks
+        # Setup embedding group (to exchange gradients between
+        # first and last stages).
+        if len(ranks) > 1:
+            embedding_ranks = [ranks[0], ranks[-1]]
+        else:
+            embedding_ranks = ranks
+        group = torch.distributed.new_group(embedding_ranks)
+        if rank in embedding_ranks:
+            _EMBEDDING_GROUP = group
+        if rank in ranks:
+            _EMBEDDING_GLOBAL_RANKS = embedding_ranks
+
+def model_parallel_is_initialized():
+    """Check if model and data parallel groups are initialized."""
+    if _TENSOR_MODEL_PARALLEL_GROUP is None or _PIPELINE_MODEL_PARALLEL_GROUP is None or _DATA_PARALLEL_GROUP is None:
+        return False
+    return True
+
+
+def get_model_parallel_group():
+    """Get the model parallel group the caller rank belongs to."""
+    assert _MODEL_PARALLEL_GROUP is not None, "model parallel group is not initialized"
+    return _MODEL_PARALLEL_GROUP
+
+
+def get_tensor_model_parallel_group():
+    """Get the tensor model parallel group the caller rank belongs to."""
+    assert _TENSOR_MODEL_PARALLEL_GROUP is not None, "intra_layer_model parallel group is not initialized"
+    return _TENSOR_MODEL_PARALLEL_GROUP
+
+
+def get_pipeline_model_parallel_group():
+    """Get the pipeline model parallel group the caller rank belongs to."""
+    assert _PIPELINE_MODEL_PARALLEL_GROUP is not None, "pipeline_model parallel group is not initialized"
+    return _PIPELINE_MODEL_PARALLEL_GROUP
+
+
+def get_data_parallel_group():
+    """Get the data parallel group the caller rank belongs to."""
+    assert _DATA_PARALLEL_GROUP is not None, "data parallel group is not initialized"
+    return _DATA_PARALLEL_GROUP
+
+
+def get_embedding_group():
+    """Get the embedding group the caller rank belongs to."""
+    assert _EMBEDDING_GROUP is not None, "embedding group is not initialized"
+    return _EMBEDDING_GROUP
+
+
+def is_rank_in_embedding_group(ignore_virtual=False):
+    """Return true if current rank is in embedding group, False otherwise."""
+    rank = torch.distributed.get_rank()
+    global _EMBEDDING_GLOBAL_RANKS
+    if ignore_virtual:
+        return rank in _EMBEDDING_GLOBAL_RANKS
+    if rank in _EMBEDDING_GLOBAL_RANKS:
+        if rank == _EMBEDDING_GLOBAL_RANKS[0]:
+            return is_pipeline_first_stage(ignore_virtual=False)
+        elif rank == _EMBEDDING_GLOBAL_RANKS[-1]:
+            return is_pipeline_last_stage(ignore_virtual=False)
+        else:
+            return True
+    return False
+
+
+def set_tensor_model_parallel_world_size(world_size):
+    """Set the tensor model parallel size"""
+    global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
+    _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = world_size
+
+
+def set_pipeline_model_parallel_world_size(world_size):
+    """Set the pipeline model parallel size"""
+    global _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
+    _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = world_size
+
+
+def get_tensor_model_parallel_world_size():
+    """Return world size for the tensor model parallel group."""
+    global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
+    if _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE is not None:
+        return _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
+    return torch.distributed.get_world_size(group=get_tensor_model_parallel_group())
+
+
+def get_pipeline_model_parallel_world_size():
+    """Return world size for the pipeline model parallel group."""
+    global _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
+    if _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE is not None:
+        return _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
+    return torch.distributed.get_world_size(group=get_pipeline_model_parallel_group())
+
+
+def set_tensor_model_parallel_rank(rank):
+    """Set tensor model parallel rank."""
+    global _MPU_TENSOR_MODEL_PARALLEL_RANK
+    _MPU_TENSOR_MODEL_PARALLEL_RANK = rank
+
+
+def set_pipeline_model_parallel_rank(rank):
+    """Set pipeline model parallel rank."""
+    global _MPU_PIPELINE_MODEL_PARALLEL_RANK
+    _MPU_PIPELINE_MODEL_PARALLEL_RANK = rank
+
+
+def get_tensor_model_parallel_rank():
+    """Return my rank for the tensor model parallel group."""
+    global _MPU_TENSOR_MODEL_PARALLEL_RANK
+    if _MPU_TENSOR_MODEL_PARALLEL_RANK is not None:
+        return _MPU_TENSOR_MODEL_PARALLEL_RANK
+    return torch.distributed.get_rank(group=get_tensor_model_parallel_group())
+
+
+def get_pipeline_model_parallel_rank():
+    """Return my rank for the pipeline model parallel group."""
+    global _MPU_PIPELINE_MODEL_PARALLEL_RANK
+    if _MPU_PIPELINE_MODEL_PARALLEL_RANK is not None:
+        return _MPU_PIPELINE_MODEL_PARALLEL_RANK
+    return torch.distributed.get_rank(group=get_pipeline_model_parallel_group())
+
+
+def is_pipeline_first_stage(ignore_virtual=False):
+    """Return True if in the first pipeline model-parallel stage, False otherwise."""
+    if not ignore_virtual:
+        if (
+            get_virtual_pipeline_model_parallel_world_size() is not None
+            and get_virtual_pipeline_model_parallel_rank() != 0
+        ):
+            return False
+    return get_pipeline_model_parallel_rank() == 0
+
+
+def is_pipeline_last_stage(ignore_virtual=False):
+    """Return True if in the last pipeline model-parallel stage, False otherwise."""
+    if not ignore_virtual:
+        virtual_pipeline_model_parallel_world_size = get_virtual_pipeline_model_parallel_world_size()
+        if virtual_pipeline_model_parallel_world_size is not None and get_virtual_pipeline_model_parallel_rank() != (
+            virtual_pipeline_model_parallel_world_size - 1
+        ):
+            return False
+    return get_pipeline_model_parallel_rank() == (get_pipeline_model_parallel_world_size() - 1)
+
+
+def get_virtual_pipeline_model_parallel_rank():
+    """Return the virtual pipeline-parallel rank."""
+    global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
+    return _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
+
+
+def set_virtual_pipeline_model_parallel_rank(rank):
+    """Set the virtual pipeline-parallel rank."""
+    global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
+    _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = rank
+
+
+def get_virtual_pipeline_model_parallel_world_size():
+    """Return the virtual pipeline-parallel world size."""
+    global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
+    return _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
+
+
+def get_tensor_model_parallel_src_rank():
+    """Calculate the global rank corresponding to the first local rank
+    in the tensor model parallel group."""
+    global_rank = torch.distributed.get_rank()
+    local_world_size = get_tensor_model_parallel_world_size()
+    return (global_rank // local_world_size) * local_world_size
+
+
+def get_pipeline_model_parallel_first_rank():
+    assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
+    return _PIPELINE_GLOBAL_RANKS[0]
+
+
+def get_pipeline_model_parallel_last_rank():
+    assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
+    last_rank_local = get_pipeline_model_parallel_world_size() - 1
+    return _PIPELINE_GLOBAL_RANKS[last_rank_local]
+
+
+def get_pipeline_model_parallel_next_rank():
+    assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
+    rank_in_pipeline = get_pipeline_model_parallel_rank()
+    world_size = get_pipeline_model_parallel_world_size()
+    return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline + 1) % world_size]
+
+
+def get_pipeline_model_parallel_prev_rank():
+    assert _PIPELINE_GLOBAL_RANKS is not None, "Pipeline parallel group is not initialized"
+    rank_in_pipeline = get_pipeline_model_parallel_rank()
+    world_size = get_pipeline_model_parallel_world_size()
+    return _PIPELINE_GLOBAL_RANKS[(rank_in_pipeline - 1) % world_size]
+
+
+def get_data_parallel_world_size():
+    """Return world size for the data parallel group."""
+    return torch.distributed.get_world_size(group=get_data_parallel_group())
+
+
+def get_data_parallel_rank():
+    """Return my rank for the data parallel group."""
+    return torch.distributed.get_rank(group=get_data_parallel_group())
+
+
+def destroy_model_parallel():
+    """Set the groups to none."""
+    global _MODEL_PARALLEL_GROUP
+    _MODEL_PARALLEL_GROUP = None
+    global _TENSOR_MODEL_PARALLEL_GROUP
+    _TENSOR_MODEL_PARALLEL_GROUP = None
+    global _PIPELINE_MODEL_PARALLEL_GROUP
+    _PIPELINE_MODEL_PARALLEL_GROUP = None
+    global _DATA_PARALLEL_GROUP
+    _DATA_PARALLEL_GROUP = None
+    global _EMBEDDING_GROUP
+    _EMBEDDING_GROUP = None
+    global _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK
+    _VIRTUAL_PIPELINE_MODEL_PARALLEL_RANK = None
+    global _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
+    _VIRTUAL_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
+    global _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE
+    _MPU_TENSOR_MODEL_PARALLEL_WORLD_SIZE = None
+    global _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE
+    _MPU_PIPELINE_MODEL_PARALLEL_WORLD_SIZE = None
+    global _MPU_TENSOR_MODEL_PARALLEL_RANK
+    _MPU_TENSOR_MODEL_PARALLEL_RANK = None
+    global _MPU_PIPELINE_MODEL_PARALLEL_RANK
+    _MPU_PIPELINE_MODEL_PARALLEL_RANK = None
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/__init__.py
new file mode 100644
index 0000000000..98bb96028c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/__init__.py
@@ -0,0 +1,8 @@
+from apex.transformer.pipeline_parallel.schedules import get_forward_backward_func
+from apex.transformer.pipeline_parallel.schedules.common import build_model
+
+
+__all__ = [
+    "get_forward_backward_func",
+    "build_model",
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/_timers.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/_timers.py
new file mode 100644
index 0000000000..55d89f351b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/_timers.py
@@ -0,0 +1,83 @@
+import time
+
+import torch
+
+
+class _Timer:
+    """Timer."""
+
+    def __init__(self, name):
+        self.name_ = name
+        self.elapsed_ = 0.0
+        self.started_ = False
+        self.start_time = time.time()
+
+    def start(self):
+        """Start the timer."""
+        assert not self.started_, "timer has already been started"
+        torch.cuda.synchronize()
+        self.start_time = time.time()
+        self.started_ = True
+
+    def stop(self):
+        """Stop the timer."""
+        assert self.started_, "timer is not started"
+        torch.cuda.synchronize()
+        self.elapsed_ += time.time() - self.start_time
+        self.started_ = False
+
+    def reset(self):
+        """Reset timer."""
+        self.elapsed_ = 0.0
+        self.started_ = False
+
+    def elapsed(self, reset=True):
+        """Calculate the elapsed time."""
+        started_ = self.started_
+        # If the timing in progress, end it first.
+        if self.started_:
+            self.stop()
+        # Get the elapsed time.
+        elapsed_ = self.elapsed_
+        # Reset the elapsed time
+        if reset:
+            self.reset()
+        # If timing was in progress, set it back.
+        if started_:
+            self.start()
+        return elapsed_
+
+
+class _Timers:
+    """Group of timers."""
+
+    def __init__(self):
+        self.timers = {}
+
+    def __call__(self, name):
+        if name not in self.timers:
+            self.timers[name] = _Timer(name)
+        return self.timers[name]
+
+    def write(self, names, writer, iteration, normalizer=1.0, reset=False):
+        """Write timers to a tensorboard writer"""
+        # currently when using add_scalars,
+        # torch.utils.add_scalars makes each timer its own run, which
+        # polutes the runs list, so we just add each as a scalar
+        assert normalizer > 0.0
+        for name in names:
+            value = self.timers[name].elapsed(reset=reset) / normalizer
+            writer.add_scalar(name + "-time", value, iteration)
+
+    def log(self, names, normalizer=1.0, reset=True):
+        """Log a group of timers."""
+        assert normalizer > 0.0
+        string = "time (ms)"
+        for name in names:
+            elapsed_time = self.timers[name].elapsed(reset=reset) * 1000.0 / normalizer
+            string += " | {}: {:.2f}".format(name, elapsed_time)
+        if torch.distributed.is_initialized():
+            if torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1):
+                print(string, flush=True)
+        else:
+            print(string, flush=True)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/p2p_communication.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/p2p_communication.py
new file mode 100644
index 0000000000..cb5abf8c43
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/p2p_communication.py
@@ -0,0 +1,398 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+from functools import reduce
+import operator
+from typing import Union, Optional, Tuple
+
+import torch
+
+from apex._autocast_utils import _get_current_dtype
+from apex.transformer import parallel_state
+from apex.transformer.utils import split_tensor_into_1d_equal_chunks
+from apex.transformer.utils import gather_split_1d_tensor
+from apex.transformer.pipeline_parallel.utils import Shape
+from apex.transformer.pipeline_parallel._timers import _Timers
+
+
+def _run_p2pops(
+        tensor_send_prev: Union[torch.Tensor, None],
+        tensor_send_next: Union[torch.Tensor, None],
+        tensor_recv_prev: Union[torch.Tensor, None],
+        tensor_recv_next: Union[torch.Tensor, None],
+):
+    ops = []
+    if tensor_send_prev is not None:
+        send_prev_op = torch.distributed.P2POp(
+            torch.distributed.isend,
+            tensor_send_prev,
+            parallel_state.get_pipeline_model_parallel_prev_rank(),
+        )
+        ops.append(send_prev_op)
+    if tensor_recv_prev is not None:
+        recv_prev_op = torch.distributed.P2POp(
+            torch.distributed.irecv,
+            tensor_recv_prev,
+            parallel_state.get_pipeline_model_parallel_prev_rank(),
+        )
+        ops.append(recv_prev_op)
+    if tensor_send_next is not None:
+        send_next_op = torch.distributed.P2POp(
+            torch.distributed.isend,
+            tensor_send_next,
+            parallel_state.get_pipeline_model_parallel_next_rank(),
+        )
+        ops.append(send_next_op)
+    if tensor_recv_next is not None:
+        recv_next_op = torch.distributed.P2POp(
+            torch.distributed.irecv,
+            tensor_recv_next,
+            parallel_state.get_pipeline_model_parallel_next_rank(),
+        )
+        ops.append(recv_next_op)
+    if len(ops) > 0:
+        reqs = torch.distributed.batch_isend_irecv(ops)
+        for req in reqs:
+            req.wait()
+
+
+# NOTE (mkozuki): Leaving `params_dytpe` as it is for future development in PyTorch, especially APEX O2 style AMP.
+# But as of v1.10, basically all tensors are torch.float32 except for output tensors of `autocast` compatible layers.
+def _communicate(
+    tensor_send_next: Optional[torch.Tensor],
+    tensor_send_prev: Optional[torch.Tensor],
+    recv_prev: bool,
+    recv_next: bool,
+    tensor_shape: Optional[Shape] = None,
+    override_scatter_gather_tensors_in_pipeline: bool = False,
+    dtype_: torch.dtype = torch.float,
+    *,
+    scatter_gather_tensors_in_pipeline: bool = True,
+    params_dtype: Optional[torch.dtype] = None,
+    fp32_residual_connection: bool = False,
+) -> Tuple[Union[torch.Tensor, None], Union[torch.Tensor, None]]:
+    """Base function for communication of tensors between stages.
+
+    Args:
+        tensor_send_next: tensor to send to next rank (no tensor sent if set to None).
+        tensor_send_prev: tensor to send to prev rank (no tensor sent if set to None).
+        recv_prev: boolean for whether tensor should be received from previous rank.
+        recv_next: boolean for whether tensor should be received from next rank.
+        tensor_shape: optional, use when the input sequence contains less tokens than the default sequence length
+        override_scatter_gather_tensors_in_pipeline:
+            optional, this is used when tensor_shape is provided to override scatter gather tensors
+        dtype_: This is used when tensor_shape is provided and what is the type of tensor_shape
+
+    Keyword args:
+        scatter_gather_tensors_in_pipeline: Optional. If :obj:`True`, use scatter/gather to optimize communication of tensors.
+        params_dtype: Optional and legacy. Defaults to torch.float. If you manually call `.half()` or `.bfloat16()` on
+            your model deliberately, pass this argument.
+        fp32_residual_connection: Optional. If :obj:`True`, move residual connections to fp32.
+
+    Returns:
+        tuple containing
+
+        - tensor_recv_prev: `torch.Tensor` if `recv_prev` is :obj:`True`, `None` otherwise.
+        - tensor_recv_next: `torch.Tensor` if `recv_next` is :obj:`True`, `None` otherwise.
+    """
+    # Create placeholder tensors for receive in forward and backward directions if needed.
+    tensor_recv_prev = None
+    tensor_recv_next = None
+    if tensor_shape is None:
+        # In megatron, `tensor_shape` is set to `(args.seq_length, args.micro_batch_size, args.hidden_size)`
+        raise RuntimeError(
+            "`tensor_shape` must be specified. Common `tensor_shape` is `(seq_length, micro_batch_size, hidden_size)`")
+    if not override_scatter_gather_tensors_in_pipeline and scatter_gather_tensors_in_pipeline:
+        tensor_chunk_shape = (reduce(operator.mul, tensor_shape, 1) // parallel_state.get_tensor_model_parallel_world_size(),)
+    else:
+        tensor_chunk_shape = tensor_shape
+    dtype = params_dtype or torch.float
+    if fp32_residual_connection:
+        dtype = torch.float
+
+    requires_grad = True
+    if dtype_ is not None:
+        dtype = dtype_
+        requires_grad = False
+
+    if recv_prev:
+        tensor_recv_prev = torch.empty(
+            tensor_chunk_shape,
+            requires_grad=requires_grad,
+            device=torch.cuda.current_device(),
+            dtype=dtype,
+        )
+    if recv_next:
+        tensor_recv_next = torch.empty(
+            tensor_chunk_shape,
+            requires_grad=requires_grad,
+            device=torch.cuda.current_device(),
+            dtype=dtype,
+        )
+
+    # Split tensor into smaller chunks if using scatter-gather optimization.
+    if not override_scatter_gather_tensors_in_pipeline and scatter_gather_tensors_in_pipeline:
+        if tensor_send_next is not None:
+            tensor_send_next = split_tensor_into_1d_equal_chunks(tensor_send_next)
+
+        if tensor_send_prev is not None:
+            tensor_send_prev = split_tensor_into_1d_equal_chunks(tensor_send_prev)
+
+    # Send tensors in both the forward and backward directions as appropriate.
+    _run_p2pops(tensor_send_prev, tensor_send_next, tensor_recv_prev, tensor_recv_next)
+    # To protect against race condition when using batch_isend_irecv().
+    torch.cuda.synchronize()
+
+    # If using scatter-gather optimization, gather smaller chunks.
+    if not override_scatter_gather_tensors_in_pipeline and scatter_gather_tensors_in_pipeline:
+        if recv_prev:
+            tensor_recv_prev = (
+                gather_split_1d_tensor(tensor_recv_prev)
+                .view(tensor_shape)
+                .requires_grad_()
+            )
+
+        if recv_next:
+            tensor_recv_next = (
+                gather_split_1d_tensor(tensor_recv_next)
+                .view(tensor_shape)
+                .requires_grad_()
+            )
+
+    return tensor_recv_prev, tensor_recv_next
+
+
+def recv_forward(
+        tensor_shape: Shape,
+        override_scatter_gather_tensors_in_pipeline: bool = False,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        timers: _Timers = None,
+) -> torch.Tensor:
+    """Receive tensor from previous rank in pipeline (forward receive)."""
+    if parallel_state.is_pipeline_first_stage():
+        return None
+    if timers is not None:
+        timers("forward-recv").start()
+    input_tensor, _ = _communicate(
+        tensor_send_next=None,
+        tensor_send_prev=None,
+        recv_prev=True,
+        recv_next=False,
+        tensor_shape=tensor_shape,
+        override_scatter_gather_tensors_in_pipeline=override_scatter_gather_tensors_in_pipeline,
+        dtype_=_get_current_dtype(dtype),
+    )
+    if timers is not None:
+        timers("forward-recv").stop()
+    return input_tensor
+
+
+def recv_backward(
+        tensor_shape: Shape = None,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        timers: _Timers = None,
+):
+    """Receive tensor from next rank in pipeline (backward receive)."""
+    if parallel_state.is_pipeline_last_stage():
+        return None
+    if timers is not None:
+        timers("backward-recv").start()
+    _, output_tensor_grad = _communicate(
+        tensor_send_next=None,
+        tensor_send_prev=None,
+        recv_prev=False,
+        recv_next=True,
+        tensor_shape=tensor_shape,
+        dtype_=_get_current_dtype(dtype),
+    )
+    if timers is not None:
+        timers("backward-recv").stop()
+    return output_tensor_grad
+
+
+def send_forward(
+        output_tensor: torch.Tensor,
+        override_scatter_gather_tensors_in_pipeline: bool = False,
+        tensor_shape: Shape = None,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        timers: _Timers = None,
+) -> None:
+    """Send tensor to next rank in pipeline (forward send)."""
+    if parallel_state.is_pipeline_last_stage():
+        return
+    if timers is not None:
+        timers("forward-send").start()
+    _communicate(
+        tensor_send_next=output_tensor,
+        tensor_send_prev=None,
+        recv_prev=False,
+        recv_next=False,
+        override_scatter_gather_tensors_in_pipeline=override_scatter_gather_tensors_in_pipeline,
+        tensor_shape=tensor_shape,
+        dtype_=_get_current_dtype(dtype),
+    )
+    if timers is not None:
+        timers("forward-send").stop()
+
+
+def send_backward(
+        input_tensor_grad: torch.Tensor,
+        tensor_shape: Shape,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        timers: _Timers = None,
+) -> None:
+    """Send tensor to previous rank in pipeline (backward send)."""
+    if parallel_state.is_pipeline_first_stage():
+        return
+    if timers is not None:
+        timers("backward-send").start()
+    _communicate(
+        tensor_send_next=None,
+        tensor_send_prev=input_tensor_grad,
+        recv_prev=False,
+        recv_next=False,
+        tensor_shape=tensor_shape,
+        dtype_=_get_current_dtype(dtype),
+    )
+    if timers is not None:
+        timers("backward-send").stop()
+
+
+def send_forward_recv_backward(
+        output_tensor: torch.Tensor,
+        tensor_shape: Shape,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        timers: _Timers = None,
+) -> None:
+    """Batched send and recv with next rank in pipeline."""
+    if parallel_state.is_pipeline_last_stage():
+        return None
+    if timers is not None:
+        timers("forward-send-backward-recv").start()
+    _, output_tensor_grad = _communicate(
+        tensor_send_next=output_tensor,
+        tensor_send_prev=None,
+        recv_prev=False,
+        recv_next=True,
+        tensor_shape=tensor_shape,
+        dtype_=_get_current_dtype(dtype),
+    )
+    if timers is not None:
+        timers("forward-send-backward-recv").stop()
+    return output_tensor_grad
+
+
+def send_backward_recv_forward(
+        input_tensor_grad: torch.Tensor,
+        tensor_shape: Shape,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        timers: _Timers = None,
+) -> torch.Tensor:
+    """Batched send and recv with previous rank in pipeline."""
+    if parallel_state.is_pipeline_first_stage():
+        return None
+    if timers is not None:
+        timers("backward-send-forward-recv").start()
+    input_tensor, _ = _communicate(
+        tensor_send_next=None,
+        tensor_send_prev=input_tensor_grad,
+        recv_prev=True,
+        recv_next=False,
+        tensor_shape=tensor_shape,
+        dtype_=_get_current_dtype(dtype),
+    )
+    if timers is not None:
+        timers("backward-send-forward-recv").stop()
+    return input_tensor
+
+
+def send_forward_recv_forward(
+        output_tensor: torch.Tensor,
+        recv_prev: bool,
+        tensor_shape: Shape,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        timers: _Timers = None,
+) -> torch.Tensor:
+    """Batched recv from previous rank and send to next rank in pipeline."""
+    if timers is not None:
+        timers("forward-send-forward-recv").start()
+    input_tensor, _ = _communicate(
+        tensor_send_next=output_tensor,
+        tensor_send_prev=None,
+        recv_prev=recv_prev,
+        recv_next=False,
+        tensor_shape=tensor_shape,
+        dtype_=_get_current_dtype(dtype),
+    )
+    if timers is not None:
+        timers("forward-send-forward-recv").stop()
+    return input_tensor
+
+
+def send_backward_recv_backward(
+        input_tensor_grad: torch.Tensor,
+        recv_next: bool,
+        tensor_shape: Shape,
+        *,
+        dtype: torch.dtype = torch.float,
+        timers: _Timers = None,
+) -> torch.Tensor:
+    """Batched recv from next rank and send to previous rank in pipeline."""
+    if timers is not None:
+        timers("backward-send-backward-recv").start()
+    _, output_tensor_grad = _communicate(
+        tensor_send_next=None,
+        tensor_send_prev=input_tensor_grad,
+        recv_prev=False,
+        recv_next=recv_next,
+        tensor_shape=tensor_shape,
+        dtype_=_get_current_dtype(dtype),
+    )
+    if timers is not None:
+        timers("backward-send-backward-recv").stop()
+    return output_tensor_grad
+
+
+def send_forward_backward_recv_forward_backward(
+        output_tensor: torch.Tensor,
+        input_tensor_grad: torch.Tensor,
+        recv_prev: bool,
+        recv_next: bool,
+        tensor_shape: Shape,
+        *,
+        dtype: Optional[torch.dtype] = None,
+        timers: _Timers = None,
+):
+    """Batched send and recv with previous and next ranks in pipeline."""
+    if timers is not None:
+        timers("forward-backward-send-forward-backward-recv").start()
+    input_tensor, output_tensor_grad = _communicate(
+        tensor_send_next=output_tensor,
+        tensor_send_prev=input_tensor_grad,
+        recv_prev=recv_prev,
+        recv_next=recv_next,
+        tensor_shape=tensor_shape,
+        dtype_=_get_current_dtype(dtype),
+    )
+    if timers is not None:
+        timers("forward-backward-send-forward-backward-recv").stop()
+    return input_tensor, output_tensor_grad
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/__init__.py
new file mode 100644
index 0000000000..61afb361fd
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/__init__.py
@@ -0,0 +1,39 @@
+import warnings
+
+from apex.transformer import parallel_state
+from apex.transformer.pipeline_parallel.utils import get_num_microbatches
+from apex.transformer.pipeline_parallel.schedules.fwd_bwd_no_pipelining import forward_backward_no_pipelining
+from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_with_interleaving import _forward_backward_pipelining_with_interleaving
+from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_without_interleaving import (
+    forward_backward_pipelining_without_interleaving,
+)
+
+
+class ExperimentalWarning(Warning):
+    pass
+
+
+def get_forward_backward_func(
+    virtual_pipeline_model_parallel_size, pipeline_model_parallel_size,
+):
+    if parallel_state.get_pipeline_model_parallel_world_size() > 1:
+        if virtual_pipeline_model_parallel_size is not None:
+            if get_num_microbatches() % pipeline_model_parallel_size != 0:
+                msg = "number of microbatches is not divisible by pipeline-parallel size when using interleaved schedule"
+                raise RuntimeError(msg)
+            warnings.warn(
+                "Pipeline Model Parallel with interleaving scheduling is experimental. "
+                f"To use Pipeline Parallel without interleaving, set `virtual_pipeline_model_parallel_size` to `None`: {virtual_pipeline_model_parallel_size}",
+                ExperimentalWarning
+            )
+            forward_backward_func = _forward_backward_pipelining_with_interleaving
+        else:
+            forward_backward_func = forward_backward_pipelining_without_interleaving
+    else:
+        forward_backward_func = forward_backward_no_pipelining
+    return forward_backward_func
+
+
+__all__ = [
+    "get_forward_backward_func",
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/common.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/common.py
new file mode 100644
index 0000000000..ad599b510f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/common.py
@@ -0,0 +1,218 @@
+# NOTE (mkozuki): For simplicity, tentatively `timers` related operations are commented out.
+from typing import Any, Callable, Dict, List, Tuple, Union, Optional
+
+import torch
+
+from apex.transformer import parallel_state
+from apex.transformer.pipeline_parallel.utils import get_num_microbatches
+from apex.transformer.pipeline_parallel.utils import listify_model
+from apex.transformer.pipeline_parallel.utils import unwrap_model
+from apex.transformer.tensor_parallel.layers import set_defaults_if_not_set_tensor_model_parallel_attributes
+
+
+Batch = Union[torch.Tensor, List[torch.Tensor], Tuple[torch.Tensor, ...]]
+LossFunc = Callable[[torch.Tensor], torch.Tensor]
+FwdStepFunc = Callable[[Batch, torch.nn.Module], Tuple[torch.Tensor, LossFunc]]
+
+
+def build_model(
+        model_provider_func: Callable[[Any, Dict[str, Any]], torch.nn.Module],
+        wrap_with_ddp: bool = True,
+        virtual_pipeline_model_parallel_size: Optional[int] = None,
+        *args,
+        **kwargs
+) -> List[torch.nn.Module]:
+    """Build the model satisfying pipeline model parallel requirements.
+
+    This function sets `pre_process` and `post_process` to `**kwargs` and pass `*args` and `**kwargs` to
+    `model_provider_func`.
+
+    Args:
+        model_provider_func: A function which takes `*args` and `**kwargs` and returns a `nn.Module`.
+        wrap_with_ddp: If :obj:`True`, wrap the instantiated model
+            with `torch.nn.parallel.distributed.DistributedDataParallel`, a.k.a. `DDP`.
+        virtual_pipeline_model_parallel_size: Specify when using interleaving scheduling pipeline model parallel.
+        *args: arguments for model provider func
+        **kwargs: Keyword arguments for model provider func
+
+    Returns:
+        a list of `nn.Module`(s). If `virtual_pipeline_model_parallel_size` is not None,
+        the list has multiple models, otherwise one.
+    """
+    if (
+            parallel_state.get_pipeline_model_parallel_world_size() > 1 and
+            virtual_pipeline_model_parallel_size is not None
+    ):
+        model = []
+        for i in range(virtual_pipeline_model_parallel_size):
+            cur_args = args
+            cur_kwargs = kwargs
+            parallel_state.set_virtual_pipeline_model_parallel_rank(i)
+            # Set pre_process and post_process only after virtual rank is set.
+            pre_process = parallel_state.is_pipeline_first_stage()
+            post_process = parallel_state.is_pipeline_last_stage()
+            cur_kwargs.update({
+                "pre_process": pre_process,
+                "post_process": post_process,
+            })
+            this_model = model_provider_func(*cur_args, **cur_kwargs)
+            model.append(this_model)
+    else:
+        cur_args = args
+        cur_kwargs = kwargs
+        pre_process = parallel_state.is_pipeline_first_stage()
+        post_process = parallel_state.is_pipeline_last_stage()
+        cur_kwargs.update({
+            "pre_process": pre_process,
+            "post_process": post_process,
+        })
+        model = model_provider_func(*cur_args, **cur_kwargs)
+
+    if not isinstance(model, list):
+        model = [model]
+
+    # Set tensor model parallel attributes if not set.
+    # Only parameters that are already tensor model parallel have these
+    # attributes set for them. We should make sure the default attributes
+    # are set for all params so the optimizer can use them.
+    for model_module in model:
+        for param in model_module.parameters():
+            set_defaults_if_not_set_tensor_model_parallel_attributes(param)
+
+    # Print number of parameters.
+    if parallel_state.get_data_parallel_rank() == 0:
+        msg = " > number of parameters on (tensor, pipeline) model parallel rank ({}, {}): {}".format(
+            parallel_state.get_tensor_model_parallel_rank(),
+            parallel_state.get_pipeline_model_parallel_rank(),
+            sum([sum([p.nelement() for p in model_module.parameters()]) for model_module in model])
+        )
+        print(msg, flush=True)
+
+    # GPU allocation.
+    for model_module in model:
+        model_module.cuda(torch.cuda.current_device())
+
+    if wrap_with_ddp:
+        i = torch.cuda.current_device()
+        model = [
+            torch.nn.parallel.distributed.DistributedDataParallel(
+                model_module,
+                device_ids=[i],
+                output_device=i,
+                process_group=parallel_state.get_data_parallel_group(),
+            )
+            for model_module in model
+        ]
+    return model
+
+
+def _get_params_for_weight_decay_optimization(
+        model: Union[torch.nn.Module, List[torch.nn.Module]],
+) -> Dict[str, torch.nn.Parameter]:
+    """Divide params into with-weight-decay and without-weight-decay groups.
+    Layernorms and biases will have no weight decay but the rest will.
+    """
+    modules = listify_model(model)
+    from apex.normalization.fused_layer_norm import FusedLayerNorm  # NOQA
+    weight_decay_params = {'params': []}
+    no_weight_decay_params = {'params': [], 'weight_decay': 0.0}
+    for module in modules:
+        for module_ in module.modules():
+            if isinstance(module_, FusedLayerNorm):
+                no_weight_decay_params['params'].extend(
+                    [p for p in list(module_._parameters.values())
+                     if p is not None])
+            else:
+                weight_decay_params['params'].extend(
+                    [p for n, p in list(module_._parameters.items())
+                     if p is not None and n != 'bias'])
+                no_weight_decay_params['params'].extend(
+                    [p for n, p in list(module_._parameters.items())
+                     if p is not None and n == 'bias'])
+
+    return weight_decay_params, no_weight_decay_params
+
+
+def forward_step(
+        forward_step_func: FwdStepFunc,
+        batch: Batch,
+        model: torch.nn.Module,
+        input_tensor: Optional[torch.Tensor],
+        losses_reduced: List[torch.Tensor],
+):
+    """Forward step for passed-in model.
+
+    If first stage, input tensor is obtained from data_iterator, otherwise
+    passed-in input_tensor is used.
+
+    Returns output tensor.
+
+    Args:
+        forward_step_func: Model specific function. This takes a minibatch and model as its arguments and
+            returns the model's output and the loss function.
+        batch: minibatch
+        model: unwrappable model
+        input_tensor:
+        losses_reduced:
+
+    Returns:
+        output_tensor
+    """
+    # timers = get_timers()
+    # timers("forward-compute").start()
+    unwrapped_model = unwrap_model(model)
+    # NOTE (mkozuki): The passed `model` is expected to implement `set_input_tensor`.
+    # See https://github.com/NVIDIA/Megatron-LM/blob/5ac5571ba0265af4c491ee0af1508ca7589450c6/megatron/model/transformer.py#L679  # NOQA
+    # for the details of `set_input_tensor`.
+    unwrapped_model.set_input_tensor(input_tensor)
+    output_tensor, loss_func = forward_step_func(batch, model)
+    # print(f"forward_step| pipeline rank: {parallel_state.get_pipeline_model_parallel_rank()} is_pipeline_last_stage?: {parallel_state.is_pipeline_last_stage()}")
+    if parallel_state.is_pipeline_last_stage():
+        output_tensor = loss_func(output_tensor)
+        loss, loss_reduced = output_tensor
+        output_tensor = loss / get_num_microbatches()
+        losses_reduced.append(loss_reduced)
+    # timers("forward-compute").stop()
+
+    return output_tensor
+
+
+def backward_step(
+        input_tensor: Optional[torch.Tensor],
+        output_tensor: torch.Tensor,
+        output_tensor_grad: Optional[torch.Tensor],
+) -> Optional[torch.Tensor]:
+    """Backward step through passed-in output tensor.
+
+    If last stage, output_tensor_grad is None, otherwise gradient of loss
+    with respect to stage's output tensor.
+
+    Returns gradient of loss with respect to input tensor (None if first
+    stage).
+
+    Args:
+        input_tensor:
+        output_tensor:
+        output_tensor_grad:
+    Returns:
+        input_tensor_grad
+    """
+
+    # timers = get_timers()
+    # timers("backward-compute").start()
+    # Retain the grad on the input_tensor.
+
+    # if parallel_state.get_pipeline_model_parallel_rank() == 0:
+    #     print(f"{input_tensor}, {output_tensor}, {output_tensor_grad}")
+    if input_tensor is not None:
+        input_tensor.retain_grad()
+    # Backward pass.
+    # if output_tensor_grad is None:
+    #     output_tensor = optimizer.scale_loss(output_tensor)
+    torch.autograd.backward(output_tensor, grad_tensors=output_tensor_grad)
+    input_tensor_grad = None
+    if input_tensor is not None:
+        input_tensor_grad = input_tensor.grad
+    # timers("backward-compute").stop()
+
+    return input_tensor_grad
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py
new file mode 100644
index 0000000000..754e6fd1a4
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_no_pipelining.py
@@ -0,0 +1,81 @@
+from contextlib import contextmanager
+from typing import List, Union
+
+import torch
+
+from apex.transformer.pipeline_parallel.utils import listify_model
+from apex.transformer.pipeline_parallel.utils import get_num_microbatches
+from apex.transformer.pipeline_parallel.utils import get_kth_microbatch
+from apex.transformer.pipeline_parallel.schedules.common import Batch, FwdStepFunc
+from apex.transformer.pipeline_parallel.schedules.common import forward_step
+from apex.transformer.pipeline_parallel.schedules.common import backward_step
+
+
+@contextmanager
+def placeholder_handler():
+    try:
+        yield
+    finally:
+        pass
+
+
+# TODO (mkozuki): Confirm this will be used or not.
+# TODO (mkozuki): Fix if necessary. Currently I'm seeing failure if `not forward_only` and
+#   the last `backward_step` seems to fail. However, note the possibility of my test script is wrong.
+def forward_backward_no_pipelining(
+        forward_step_func: FwdStepFunc,
+        batch: Batch,
+        model: Union[torch.nn.Module, List[torch.nn.Module]],
+        *,
+        forward_only: bool,
+        **kwargs,
+):
+    """Run forward and backward passes with no pipeline parallelism (no inter-stage communication).
+
+    This pipeline parallel scheduling handles the last microbatch differently to synchronize gradients.
+
+    Args:
+        forward_step_func: A function which takes a minibatch and model as its arguments and
+            returns model's forward output and the loss function.
+            The loss function is supposed to take one `torch.Tensor` and
+            return a `torch.Tensor` of loss and a dictionary of `str` and `torch.Tensor`.
+        batch: A List of torch.Tensors
+        model: A `torch.nn.Module` or a list of `torch.nn.Module`.
+
+    Keyword args:
+        forward_only:
+        **kwargs: Added to handle `tensor_shape` which has no effect on this function.
+
+    Returns:
+        a list of dictionaries of loss `torch.Tensor`s if the last stage, empty list otherwise.
+    """
+    model = listify_model(model)
+    if len(model) != 1:
+        msg = f"`model` is expected be a `nn.Module`, but {type(model)}"
+        raise RuntimeError(msg)
+    model = model[0]
+
+    context_handler = placeholder_handler
+    if isinstance(model, torch.nn.parallel.distributed.DistributedDataParallel):
+        context_handler = model.no_sync
+
+    losses_reduced = []
+    input_tensor, output_tensor_grad = None, None
+    num_micro_batches = get_num_microbatches()
+    with context_handler():
+        for i in range(num_micro_batches - 1):
+            cur_micro_batch = get_kth_microbatch(batch, i)
+            output_tensor = forward_step(
+                forward_step_func, cur_micro_batch, model, input_tensor, losses_reduced)
+            if not forward_only:
+                backward_step(input_tensor, output_tensor, output_tensor_grad)
+
+    # Run computation for last microbatch out of context handler (want to
+    # synchronize gradients).
+    output_tensor = forward_step(
+        forward_step_func, get_kth_microbatch(batch, num_micro_batches - 1), model, input_tensor, losses_reduced
+    )
+    if not forward_only:
+        backward_step(input_tensor, output_tensor, output_tensor_grad)
+
+    return losses_reduced
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py
new file mode 100644
index 0000000000..4d3753f578
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_with_interleaving.py
@@ -0,0 +1,298 @@
+from typing import List, Union, Optional
+
+import torch
+
+from apex.transformer import parallel_state
+from apex.transformer.pipeline_parallel import p2p_communication
+from apex.transformer.pipeline_parallel.schedules.common import Batch, FwdStepFunc
+from apex.transformer.pipeline_parallel.schedules.common import backward_step
+from apex.transformer.pipeline_parallel.schedules.common import forward_step
+from apex.transformer.pipeline_parallel.utils import get_kth_microbatch
+from apex.transformer.pipeline_parallel.utils import get_num_microbatches
+from apex.transformer.utils import rank_print
+
+
+# TODO (mkozuki): Reduce cyclomatic complexity
+def _forward_backward_pipelining_with_interleaving(
+        forward_step_func: FwdStepFunc,
+        batch: List[Batch],
+        model: List[torch.nn.Module],
+        *,
+        forward_only: bool,
+        tensor_shape: Optional[Union[List[int], torch.Size]] = None,
+):
+    """Run interleaved 1F1B schedule with communication between pipeline stages as needed.
+
+    This function assumes `batch` and `model` is a list of `Batch`'s and a list of `torch.nn.Module`, respectively.
+    This means that model is split into model chunks.
+
+    This pipeline parallel scheduling consists of three steps:
+        1. warmup
+        2. 1F1B a.k.a. steady state
+        3. cooldown
+    Note that if `forward_only` this scheduling consists of only warmup phase.
+
+    Args:
+        forward_step_func: A function which takes a minibatch and model as its arguments and
+            returns model's forward output and the loss function.
+            The loss function is supposed to take one `torch.Tensor` and
+            return a `torch.Tensor` of loss and a dictionary of `str` and `torch.Tensor`.
+        batch: A minibatch, i.e., a list of `torch.Tensor`'s.
+        model: A `torch.nn.Module` or a list of `torch.nn.Module`.
+
+    Keyword args:
+        forward_only:
+        tensor_shape: Shape of tensor.
+
+    Returns:
+        a list of loss `torch.Tensor`s if the last stage, empty list otherwise.
+    """
+    if not isinstance(model, list):
+        raise RuntimeError("`model` must be a list of `nn.Module`'s'")
+
+    num_model_chunks = len(model)
+    input_tensors = [[] for _ in range(num_model_chunks)]
+    output_tensors = [[] for _ in range(num_model_chunks)]
+    curr_iters = [0 for _ in range(num_model_chunks)]
+    losses_reduced = []
+    if not forward_only:
+        output_tensor_grads = [[] for _ in range(num_model_chunks)]
+
+    pipeline_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
+    pipeline_parallel_rank = parallel_state.get_pipeline_model_parallel_rank()
+
+    # Compute number of warmup and remaining microbatches.
+    num_microbatches = get_num_microbatches() * num_model_chunks
+    all_warmup_microbatches = False
+    if forward_only:
+        num_warmup_microbatches = num_microbatches
+    else:
+        # Run all forward passes and then all backward passes if number of
+        # microbatches is just the number of pipeline stages.
+        # Otherwise, perform (num_model_chunks-1)*pipeline_parallel_size on
+        # all workers, followed by more microbatches after depending on
+        # stage ID (more forward passes for earlier stages, later stages can
+        # immediately start with 1F1B).
+        if get_num_microbatches() == pipeline_parallel_size:
+            num_warmup_microbatches = num_microbatches
+            all_warmup_microbatches = True
+        else:
+            num_warmup_microbatches = (pipeline_parallel_size - pipeline_parallel_rank - 1) * 2
+            num_warmup_microbatches += (num_model_chunks - 1) * pipeline_parallel_size
+            num_warmup_microbatches = min(num_warmup_microbatches, num_microbatches)
+    num_microbatches_remaining = num_microbatches - num_warmup_microbatches
+
+
+    # TODO (mkozuki): Remove once debug gets done
+    # rank_print(
+    #     f"num_microbatches: {num_microbatches}, "
+    #     f"num_warmup_microbatches: {num_warmup_microbatches}, "
+    #     f"num_microbatches_remaining: {num_microbatches_remaining} -- "
+    # )
+
+    ###################################################################################################################
+    # Helper function definitions.
+    ###################################################################################################################
+    def get_model_chunk_id(microbatch_id: int, forward: bool) -> int:
+        """Helper function to get the model chunk ID given the iteration number."""
+        pipeline_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
+        microbatch_id_in_group = microbatch_id % (pipeline_parallel_size * num_model_chunks)
+        model_chunk_id = microbatch_id_in_group // pipeline_parallel_size
+        if not forward:
+            model_chunk_id = num_model_chunks - model_chunk_id - 1
+        return model_chunk_id
+
+    def forward_step_helper(microbatch_id, curr_iters):
+        """Helper method to run forward step with model split into chunks
+        (run set_virtual_pipeline_model_parallel_rank() before calling
+        forward_step())."""
+        model_chunk_id = get_model_chunk_id(microbatch_id, forward=True)
+        parallel_state.set_virtual_pipeline_model_parallel_rank(model_chunk_id)
+
+        # forward step
+        if (
+                parallel_state.is_pipeline_first_stage() and
+                len(input_tensors[model_chunk_id]) == len(output_tensors[model_chunk_id])
+        ):
+            input_tensors[model_chunk_id].append(None)
+        input_tensor = input_tensors[model_chunk_id][-1]
+        output_tensor = forward_step(
+            forward_step_func,
+            get_kth_microbatch(batch, curr_iters[model_chunk_id]),
+            model[model_chunk_id],
+            input_tensor,
+            losses_reduced,
+        )
+        curr_iters[model_chunk_id] += 1
+        output_tensors[model_chunk_id].append(output_tensor)
+
+        # if forward-only, no need to save tensors for a backward pass
+        if forward_only:
+            input_tensors[model_chunk_id].pop()
+            output_tensors[model_chunk_id].pop()
+
+        return output_tensor
+
+    def backward_step_helper(microbatch_id):
+        """Helper method to run backward step with model split into chunks
+        (run set_virtual_pipeline_model_parallel_rank() before calling
+        backward_step())."""
+        model_chunk_id = get_model_chunk_id(microbatch_id, forward=False)
+        parallel_state.set_virtual_pipeline_model_parallel_rank(model_chunk_id)
+
+        if parallel_state.is_pipeline_last_stage():
+            if len(output_tensor_grads[model_chunk_id]) == 0:
+                output_tensor_grads[model_chunk_id].append(None)
+        input_tensor = input_tensors[model_chunk_id].pop(0)
+        output_tensor = output_tensors[model_chunk_id].pop(0)
+        output_tensor_grad = output_tensor_grads[model_chunk_id].pop(0)
+        input_tensor_grad = backward_step(input_tensor, output_tensor, output_tensor_grad)
+
+        return input_tensor_grad
+
+    ###################################################################################################################
+    # Run warmup forward passes.
+    ###################################################################################################################
+    parallel_state.set_virtual_pipeline_model_parallel_rank(0)
+    input_tensors[0].append(p2p_communication.recv_forward(tensor_shape=tensor_shape))
+    for k in range(num_warmup_microbatches):
+        # rank_print(f"warmup iter: {k}")
+        output_tensor = forward_step_helper(k, curr_iters)
+
+        # Determine if tensor should be received from previous stage.
+        next_forward_model_chunk_id = get_model_chunk_id(k + 1, forward=True)
+        recv_prev = True
+        if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
+            if next_forward_model_chunk_id == 0:
+                recv_prev = False
+        if k == (num_microbatches - 1):
+            recv_prev = False
+
+        # Don't send tensor downstream if on last stage.
+        if parallel_state.is_pipeline_last_stage():
+            output_tensor = None
+
+        # rank_print(f"recv_prev: {recv_prev}")
+        # Send and receive tensors as appropriate (send tensors computed
+        # in this iteration; receive tensors for next iteration).
+        if k == (num_warmup_microbatches - 1) and not forward_only and not all_warmup_microbatches:
+            input_tensor_grad = None
+            recv_next = True
+            # rank_print(f"recv_next: {recv_next}")
+            if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+                recv_next = False
+            (
+                input_tensor,
+                output_tensor_grad,
+            ) = p2p_communication.send_forward_backward_recv_forward_backward(
+                output_tensor,
+                input_tensor_grad,
+                recv_prev=recv_prev,
+                recv_next=recv_next,
+                tensor_shape=tensor_shape,
+            )
+            output_tensor_grads[num_model_chunks - 1].append(output_tensor_grad)
+        else:
+            # rank_print("send_forward_recv_forward start")
+            input_tensor = p2p_communication.send_forward_recv_forward(output_tensor, recv_prev=recv_prev, tensor_shape=tensor_shape)
+            # rank_print("send_forward_recv_forward finish")
+        # rank_print("communication done")
+        input_tensors[next_forward_model_chunk_id].append(input_tensor)
+
+    ###################################################################################################################
+    # Run 1F1B in steady state.
+    ###################################################################################################################
+    for k in range(num_microbatches_remaining):
+        # Forward pass.
+        forward_k = k + num_warmup_microbatches
+        output_tensor = forward_step_helper(forward_k, curr_iters)
+
+        # Backward pass.
+        backward_k = k
+        input_tensor_grad = backward_step_helper(backward_k)
+
+        # Send output_tensor and input_tensor_grad, receive input_tensor
+        # and output_tensor_grad.
+
+        # Determine if current stage has anything to send in either direction,
+        # otherwise set tensor to None.
+        forward_model_chunk_id = get_model_chunk_id(forward_k, forward=True)
+        parallel_state.set_virtual_pipeline_model_parallel_rank(forward_model_chunk_id)
+        if parallel_state.is_pipeline_last_stage():
+            output_tensor = None
+
+        backward_model_chunk_id = get_model_chunk_id(backward_k, forward=False)
+        parallel_state.set_virtual_pipeline_model_parallel_rank(backward_model_chunk_id)
+        if parallel_state.is_pipeline_first_stage():
+            input_tensor_grad = None
+
+        # Determine if peers are sending, and where in data structure to put
+        # received tensors.
+        recv_prev = True
+        if parallel_state.is_pipeline_first_stage(ignore_virtual=True):
+            # First stage is ahead of last stage by (pipeline_parallel_size - 1).
+            next_forward_model_chunk_id = get_model_chunk_id(
+                forward_k - (pipeline_parallel_size - 1), forward=True
+            )
+            if next_forward_model_chunk_id == (num_model_chunks - 1):
+                recv_prev = False
+            next_forward_model_chunk_id += 1
+        else:
+            next_forward_model_chunk_id = get_model_chunk_id(forward_k + 1, forward=True)
+
+        recv_next = True
+        if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+            # Last stage is ahead of first stage by (pipeline_parallel_size - 1).
+            next_backward_model_chunk_id = get_model_chunk_id(
+                backward_k - (pipeline_parallel_size - 1), forward=False
+            )
+            if next_backward_model_chunk_id == 0:
+                recv_next = False
+            next_backward_model_chunk_id -= 1
+        else:
+            next_backward_model_chunk_id = get_model_chunk_id(backward_k + 1, forward=False)
+
+        # If last iteration, don't receive; we already received one extra
+        # before the start of the for loop.
+        if k == (num_microbatches_remaining - 1):
+            recv_prev = False
+
+        # Communicate tensors.
+        (
+            input_tensor,
+            output_tensor_grad,
+        ) = p2p_communication.send_forward_backward_recv_forward_backward(
+            output_tensor,
+            input_tensor_grad,
+            recv_prev=recv_prev,
+            recv_next=recv_next,
+            tensor_shape=tensor_shape,
+        )
+
+        # Put input_tensor and output_tensor_grad in data structures in the
+        # right location.
+        if recv_prev:
+            input_tensors[next_forward_model_chunk_id].append(input_tensor)
+        if recv_next:
+            output_tensor_grads[next_backward_model_chunk_id].append(output_tensor_grad)
+
+    ###################################################################################################################
+    # Run cooldown backward passes (flush out pipeline).
+    ###################################################################################################################
+    if not forward_only:
+        if all_warmup_microbatches:
+            output_tensor_grads[num_model_chunks - 1].append(p2p_communication.recv_backward(tensor_shape=tensor_shape))
+        for k in range(num_microbatches_remaining, num_microbatches):
+            input_tensor_grad = backward_step_helper(k)
+            next_backward_model_chunk_id = get_model_chunk_id(k + 1, forward=False)
+            recv_next = True
+            if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+                if next_backward_model_chunk_id == (num_model_chunks - 1):
+                    recv_next = False
+            if k == (num_microbatches - 1):
+                recv_next = False
+            output_tensor_grads[next_backward_model_chunk_id].append(
+                p2p_communication.send_backward_recv_backward(input_tensor_grad, recv_next=recv_next, tensor_shape=tensor_shape)
+            )
+
+    return losses_reduced
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py
new file mode 100644
index 0000000000..e6bd6e3b1a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/schedules/fwd_bwd_pipelining_without_interleaving.py
@@ -0,0 +1,175 @@
+from typing import Union, List, Optional
+
+import torch
+
+from apex.transformer import parallel_state
+from apex.transformer.pipeline_parallel import p2p_communication
+from apex.transformer.pipeline_parallel.utils import get_kth_microbatch
+from apex.transformer.pipeline_parallel.utils import listify_model
+from apex.transformer.pipeline_parallel.utils import get_num_microbatches
+from apex.transformer.pipeline_parallel.schedules.common import Batch, FwdStepFunc
+from apex.transformer.pipeline_parallel.schedules.common import forward_step
+from apex.transformer.pipeline_parallel.schedules.common import backward_step
+from apex.transformer.utils import rank_print
+
+
+def forward_backward_pipelining_without_interleaving(
+        forward_step_func: FwdStepFunc,
+        batch: Batch,
+        model: Union[torch.nn.Module, List[torch.nn.Module]],
+        *,
+        forward_only: bool,
+        tensor_shape: Optional[Union[List[int], torch.Size]] = None,
+):
+    """Run non-interleaved 1F1B schedule, with communication between pipeline stages.
+
+    This pipeline parallel scheduling consists of three steps:
+        1. warmup
+        2. 1F1B a.k.a. steady state
+        3. cooldown if not forward_only
+
+    Args:
+        forward_step_func: A function which takes a minibatch and model as its arguments and
+            returns model's forward output and the loss function.
+            The loss function is supposed to take one `torch.Tensor` and
+            return a `torch.Tensor` of loss and a dictionary of `str` and `torch.Tensor`.
+        batch: A minibatch, i.e., a list of `torch.Tensor`'s.
+        model: A `torch.nn.Module` or a list of `torch.nn.Module`.
+
+    Keyword args:
+        forward_only:
+        tensor_shape: Shape of tensor. Required for P2P communication.
+
+    Returns:
+        a list of loss `torch.Tensor`s if the last stage, empty list otherwise.
+    """
+    # timers = get_timers()
+
+    model = listify_model(model)
+    if len(model) != 1:
+        msg = f"`model` is expected be a `nn.Module`, but {type(model)}"
+        raise RuntimeError(msg)
+    model = model[0]
+
+    # Compute number of warmup microbatches.
+    num_microbatches = get_num_microbatches()
+    num_warmup_microbatches = (
+        parallel_state.get_pipeline_model_parallel_world_size()
+        - parallel_state.get_pipeline_model_parallel_rank()
+        - 1
+    )
+    num_warmup_microbatches = min(num_warmup_microbatches, num_microbatches)
+    num_microbatches_remaining = num_microbatches - num_warmup_microbatches
+
+    # TODO (mkozuki): Remove once debug gets done
+    print(
+        f">>> rank: {torch.distributed.get_rank()}, "
+        f"num_microbatches: {num_microbatches}, "
+        f"num_warmup_microbatches: {num_warmup_microbatches}, "
+        f"num_microbatches_remaining: {num_microbatches_remaining} -- "
+    )
+
+    # Input, output tensors only need to be saved when doing backward passes
+    input_tensors = None
+    output_tensors = None
+    if not forward_only:
+        input_tensors = []
+        output_tensors = []
+    losses_reduced = []
+
+    ###################################################################################################################
+    # Run warmup forward passes.
+    ###################################################################################################################
+    # rank_print(f"warmup: {num_warmup_microbatches}")
+    for i in range(num_warmup_microbatches):
+        input_tensor = p2p_communication.recv_forward(tensor_shape=tensor_shape)
+        cur_microbatch = get_kth_microbatch(batch, i)
+        output_tensor = forward_step(forward_step_func, cur_microbatch, model, input_tensor, losses_reduced)
+        p2p_communication.send_forward(output_tensor, tensor_shape=tensor_shape)
+
+        if not forward_only:
+            input_tensors.append(input_tensor)
+            output_tensors.append(output_tensor)
+        # rank_print(f"warmup iter: {i + 1} / {num_warmup_microbatches}")
+    # rank_print("warmup done")
+
+    # Before running 1F1B, need to receive first forward tensor.
+    # If all microbatches are run in warmup / cooldown phase, then no need to
+    # receive this tensor here.
+    # rank_print(f"num microbatches remaining: {num_microbatches_remaining}")
+    if num_microbatches_remaining > 0:
+        # rank_print(f"recv_forward before steady state start")
+        input_tensor = p2p_communication.recv_forward(tensor_shape=tensor_shape)
+        # rank_print(f"recv_forward before steady state done")
+
+    ###################################################################################################################
+    # Run 1F1B in steady state.
+    ###################################################################################################################
+    # rank_print(f"steady: {num_microbatches_remaining} iters")
+    for i in range(num_microbatches_remaining):
+        # rank_print(f"steady: iter {i + 1} / {num_microbatches_remaining} iters")
+        # if not forward_only:
+        #     rank_print(f"len(input_tensors) = {len(input_tensors)}, len(output_tensors) = {len(output_tensors)}")
+        last_iteration = i == (num_microbatches_remaining - 1)
+
+        cur_microbatch = get_kth_microbatch(batch, i + num_warmup_microbatches)
+        output_tensor = forward_step(forward_step_func, cur_microbatch, model, input_tensor, losses_reduced)
+        if forward_only:
+            # rank_print(f"steady, no backward: `send_forward` start")
+            p2p_communication.send_forward(output_tensor, tensor_shape=tensor_shape)
+
+            if not last_iteration:
+                input_tensor = p2p_communication.recv_forward(tensor_shape=tensor_shape)
+            # rank_print(f"steady, no backward: `send_forward` finish")
+
+        else:
+            # rank_print("L.124 steady, backward: `send_forward_recv_backward` start")
+            output_tensor_grad = p2p_communication.send_forward_recv_backward(output_tensor, tensor_shape=tensor_shape)
+            # rank_print("L.124 steady, backward: `send_forward_recv_backward` finish")
+
+            # Add input_tensor and output_tensor to end of list.
+            input_tensors.append(input_tensor)
+            output_tensors.append(output_tensor)
+
+            # Pop input_tensor and output_tensor from the start of the list for the backward pass.
+            input_tensor = input_tensors.pop(0)
+            output_tensor = output_tensors.pop(0)
+
+            input_tensor_grad = backward_step(
+                input_tensor, output_tensor, output_tensor_grad
+            )
+
+            if last_iteration:
+                input_tensor = None
+                # rank_print(f"L.142 steady backward last iteration: `send_backward` start")
+                p2p_communication.send_backward(input_tensor_grad, tensor_shape=tensor_shape)
+                # rank_print(f"L.142 steady backward last iteration: `send_backward` finish")
+            else:
+                # rank_print(f"L.146 steady backward: `send_backward_recv_forward` start")
+                input_tensor = p2p_communication.send_backward_recv_forward(
+                    input_tensor_grad, tensor_shape=tensor_shape)
+                # rank_print(f"L.146 steady backward: `send_backward_recv_forward` finish")
+    # rank_print(f"steady: exit")
+    ###################################################################################################################
+    # Run cooldown backward passes.
+    ###################################################################################################################
+    if not forward_only:
+        # rank_print(f"cooldownk: {num_warmup_microbatches} iters")
+        for i in range(num_warmup_microbatches):
+            # rank_print(f"cooldown iter: {i + 1} / {num_warmup_microbatches}")
+            input_tensor = input_tensors.pop(0)
+            output_tensor = output_tensors.pop(0)
+
+            # rank_print(f"cooldown waiting for grad tensor")
+            output_tensor_grad = p2p_communication.recv_backward(tensor_shape=tensor_shape)
+
+            # rank_print(f"cooldown received grad tensor")
+            input_tensor_grad = backward_step(
+                input_tensor, output_tensor, output_tensor_grad
+            )
+
+            # rank_print(f"cooldown sending grad tensor")
+            p2p_communication.send_backward(input_tensor_grad, tensor_shape=tensor_shape)
+        # rank_print(f"cooldownk exit")
+
+    return losses_reduced
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/utils.py
new file mode 100644
index 0000000000..07c16902dc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/pipeline_parallel/utils.py
@@ -0,0 +1,317 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+"""Utilities for pipeline model parallel."""
+from typing import Optional, List, Union
+
+import torch
+from torch.nn.parallel import DistributedDataParallel
+
+from apex.multi_tensor_apply import multi_tensor_applier
+from apex.transformer import parallel_state
+from apex.transformer.microbatches import build_num_microbatches_calculator
+from apex.transformer.pipeline_parallel._timers import _Timers
+if multi_tensor_applier.available:
+    import amp_C
+
+
+_GLOBAL_ARGS = None
+_GLOBAL_NUM_MICROBATCHES_CALCULATOR = None
+_GLOBAL_TOKENIZER = None
+_GLOBAL_TENSORBOARD_WRITER = None
+_GLOBAL_AUTORESUME = None
+_GLOBAL_TIMERS = None
+
+
+Shape = Union[List[int], torch.Size]
+
+
+def listify_model(model: Union[torch.nn.Module, List[torch.nn.Module]]) -> List[torch.nn.Module]:
+    if isinstance(model, list):
+        return model
+    return [model]
+
+
+def _ensure_var_is_initialized(var, name):
+    """Make sure the input variable is not None."""
+    assert var is not None, "{} is not initialized.".format(name)
+
+
+def _ensure_var_is_not_initialized(var, name):
+    """Make sure the input variable is not None."""
+    assert var is None, "{} is already initialized.".format(name)
+
+
+def setup_microbatch_calculator(
+        rank: int,
+        rampup_batch_size: Optional[List[int]],
+        global_batch_size: int,
+        micro_batch_size: int,
+        data_parallel_size: int,
+) -> None:
+    global _GLOBAL_NUM_MICROBATCHES_CALCULATOR
+    _ensure_var_is_not_initialized(_GLOBAL_NUM_MICROBATCHES_CALCULATOR, 'num microbatches calculator')
+
+    _GLOBAL_NUM_MICROBATCHES_CALCULATOR = build_num_microbatches_calculator(
+        rank, rampup_batch_size, global_batch_size, micro_batch_size, data_parallel_size)
+
+
+def get_micro_batch_size():
+    return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.micro_batch_size
+
+
+def get_num_microbatches():
+    return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get()
+
+
+def get_current_global_batch_size():
+    return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get_current_global_batch_size()
+
+
+def update_num_microbatches(consumed_samples, consistency_check=True):
+    _GLOBAL_NUM_MICROBATCHES_CALCULATOR.update(consumed_samples, consistency_check)
+
+
+# note (mkozuki): Comment out in favor of `get_kth_microbatch`
+def _split_batch_into_microbatch(
+        batch: List[torch.Tensor],
+        *,
+        _micro_batch_size: Optional[int] = None,
+        _global_batch_size: Optional[int] = None,
+) -> List[List[torch.Tensor]]:
+    micro_batch_size = _micro_batch_size
+    global_batch_size = _global_batch_size
+    if micro_batch_size is None:
+        micro_batch_size = get_micro_batch_size()
+    if global_batch_size is None:
+        global_batch_size = get_current_global_batch_size()
+    for i in range(0, global_batch_size, micro_batch_size):
+        yield [x[i * micro_batch_size:(i + 1) * micro_batch_size] for x in batch]
+
+
+# TODO(mkozuki): Support non-tensor local minibatches?
+def get_kth_microbatch(batch: List[torch.Tensor], k: int) -> List[torch.Tensor]:
+    """Create a list of microbatches from a list of local minibatches.
+
+    This function creates a list of `k`th microbatches from a list of local minibatches.
+    `a local minibatch` consists of `global_batch_size / data_parallel_size` samples.
+    """
+    micro_batch_size = get_micro_batch_size()
+    return [x[k * micro_batch_size:(k + 1) * micro_batch_size] for x in batch]
+
+
+def get_autoresume():
+    return _GLOBAL_AUTORESUME
+
+
+def _set_timers():
+    """Initialize timers."""
+    global _GLOBAL_TIMERS
+    _ensure_var_is_not_initialized(_GLOBAL_TIMERS, "timers")
+    _GLOBAL_TIMERS = _Timers()
+
+
+def get_timers():
+    """Return timers."""
+    _ensure_var_is_initialized(_GLOBAL_TIMERS, "timers")
+    return _GLOBAL_TIMERS
+
+
+def print_rank_0(message: str) -> None:
+    """If distributed is initialized, print only on rank 0."""
+    if torch.distributed.is_initialized():
+        if torch.distributed.get_rank() == 0:
+            print(message, flush=True)
+    else:
+        print(message, flush=True)
+
+
+def is_last_rank():
+    return torch.distributed.get_rank() == (torch.distributed.get_world_size() - 1)
+
+
+def print_rank_last(message):
+    """If distributed is initialized, print only on last rank."""
+    if torch.distributed.is_initialized():
+        if is_last_rank():
+            print(message, flush=True)
+    else:
+        print(message, flush=True)
+
+
+def param_is_not_shared(param: torch.nn.Parameter) -> bool:
+    return getattr(param, "shared", False)
+
+
+def unwrap_model(model, module_instances=(DistributedDataParallel,)):
+    return_list = True
+    if not isinstance(model, list):
+        model = [model]
+        return_list = False
+    unwrapped_model = []
+    for model_module in model:
+        while isinstance(model_module, module_instances):
+            model_module = model_module.module
+        unwrapped_model.append(model_module)
+    if not return_list:
+        return unwrapped_model[0]
+    return unwrapped_model
+
+
+def calc_params_l2_norm(model: torch.nn.Module, bf16: bool):
+    """Calculate l2 norm of parameters """
+    # args = get_args()
+    if not isinstance(model, list):
+        model = [model]
+    # Remove duplicate params.
+    params_data = []
+    for model_ in model:
+        for param in model_.parameters():
+            is_not_shared = param_is_not_shared(param)
+            is_not_tp_duplicate = parallel_state.param_is_not_tensor_parallel_duplicate(param)
+            if is_not_shared and is_not_tp_duplicate:
+                if bf16:
+                    params_data.append(param.data.float())
+                else:
+                    params_data.append(param.data)
+    # Calculate norm
+    dummy_overflow_buf = torch.cuda.IntTensor([0])
+    norm, _ = multi_tensor_applier(
+        amp_C.multi_tensor_l2norm, dummy_overflow_buf, [params_data], False  # no per-parameter norm
+    )
+    norm_2 = norm * norm
+    # Sum across all model-parallel GPUs.
+    torch.distributed.all_reduce(
+        norm_2, op=torch.distributed.ReduceOp.SUM, group=parallel_state.get_model_parallel_group()
+    )
+    return norm_2.item() ** 0.5
+
+
+def average_losses_across_data_parallel_group(losses):
+    """Reduce a tensor of losses across all GPUs."""
+    averaged_losses = torch.cat([loss.clone().detach().view(1) for loss in losses])
+    torch.distributed.all_reduce(averaged_losses, group=parallel_state.get_data_parallel_group())
+    averaged_losses = averaged_losses / torch.distributed.get_world_size(
+        group=parallel_state.get_data_parallel_group()
+    )
+
+    return averaged_losses
+
+
+def report_memory(name):
+    """Simple GPU memory report."""
+    mega_bytes = 1024.0 * 1024.0
+    string = name + " memory (MB)"
+    string += " | allocated: {}".format(torch.cuda.memory_allocated() / mega_bytes)
+    string += " | max allocated: {}".format(torch.cuda.max_memory_allocated() / mega_bytes)
+    string += " | reserved: {}".format(torch.cuda.memory_reserved() / mega_bytes)
+    string += " | max reserved: {}".format(torch.cuda.max_memory_reserved() / mega_bytes)
+    if parallel_state.get_data_parallel_rank() == 0:
+        print("[Rank {}] {}".format(torch.distributed.get_rank(), string), flush=True)
+
+
+def print_params_min_max_norm(optimizer, iteration):
+    """Print min, max, and norm of all parameters."""
+    index = 0
+    rank = torch.distributed.get_rank()
+    string = "iteration, rank, index, tensor-model-parallel, min, max, norm\n"
+    optimizer_ = optimizer.optimizer
+    for param_group in optimizer_.param_groups:
+        for param in param_group["params"]:
+            index += 1
+            min_ = param.data.min()
+            max_ = param.data.max()
+            norm = torch.linalg.norm(param.data)
+            string += "{:7d}, {:4d}, {:4d}, {:2d}, ".format(
+                iteration, rank, index, int(param.tensor_model_parallel)
+            )
+            string += "{:.6E}, {:.6E}, {:.6E}\n".format(min_, max_, norm)
+    print(string, flush=True)
+
+
+# NOTE (mkozuki): APEX doesn't have anything equivalent for
+# `_GLOBAL_ADLR_AUTORESUME` like Megatron-LM.
+# def check_adlr_autoresume_termination(iteration, model, optimizer, lr_scheduler, save: bool):
+#     """Check for autoresume signal and exit if it is received."""
+#     from apex.ppu.checkpointing import save_checkpoint
+#
+#     autoresume = get_adlr_autoresume()
+#     # Add barrier to ensure consistency.
+#     torch.distributed.barrier()
+#     if autoresume.termination_requested():
+#         if save:
+#             save_checkpoint(iteration, model, optimizer, lr_scheduler)
+#         print_rank_0(">>> autoresume termination request found!")
+#         if torch.distributed.get_rank() == 0:
+#             autoresume.request_resume()
+#         print_rank_0(">>> training terminated. Returning")
+#         sys.exit(0)
+
+
+def get_ltor_masks_and_position_ids(
+    data, eod_token, reset_position_ids, reset_attention_mask, eod_mask_loss
+):
+    """Build masks and position id for left to right model."""
+
+    # Extract batch size and sequence length.
+    micro_batch_size, seq_length = data.size()
+
+    # Attention mask (lower triangular).
+    if reset_attention_mask:
+        att_mask_batch = micro_batch_size
+    else:
+        att_mask_batch = 1
+    attention_mask = torch.tril(
+        torch.ones((att_mask_batch, seq_length, seq_length), device=data.device)
+    ).view(att_mask_batch, 1, seq_length, seq_length)
+
+    # Loss mask.
+    loss_mask = torch.ones(data.size(), dtype=torch.float, device=data.device)
+    if eod_mask_loss:
+        loss_mask[data == eod_token] = 0.0
+
+    # Position ids.
+    position_ids = torch.arange(seq_length, dtype=torch.long, device=data.device)
+    position_ids = position_ids.unsqueeze(0).expand_as(data)
+    # We need to clone as the ids will be modifed based on batch index.
+    if reset_position_ids:
+        position_ids = position_ids.clone()
+
+    if reset_position_ids or reset_attention_mask:
+        # Loop through the batches:
+        for b in range(micro_batch_size):
+
+            # Find indecies where EOD token is.
+            eod_index = position_ids[b, data[b] == eod_token]
+            # Detach indecies from positions if going to modify positions.
+            if reset_position_ids:
+                eod_index = eod_index.clone()
+
+            # Loop through EOD indecies:
+            prev_index = 0
+            for j in range(eod_index.size()[0]):
+                i = eod_index[j]
+                # Mask attention loss.
+                if reset_attention_mask:
+                    attention_mask[b, 0, (i + 1) :, : (i + 1)] = 0
+                # Reset positions.
+                if reset_position_ids:
+                    position_ids[b, (i + 1) :] -= i + 1 - prev_index
+                    prev_index = i + 1
+
+    # Convert attention mask to binary:
+    attention_mask = attention_mask < 0.5
+
+    return attention_mask, loss_mask, position_ids
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/__init__.py
new file mode 100644
index 0000000000..f4e02b1121
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/__init__.py
@@ -0,0 +1,74 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+"""Model parallel utility interface."""
+
+from apex.transformer.tensor_parallel.cross_entropy import vocab_parallel_cross_entropy
+
+from apex.transformer.tensor_parallel.data import broadcast_data
+
+from apex.transformer.tensor_parallel.layers import (
+    ColumnParallelLinear,
+    RowParallelLinear,
+    VocabParallelEmbedding,
+    set_tensor_model_parallel_attributes,
+    set_defaults_if_not_set_tensor_model_parallel_attributes,
+    copy_tensor_model_parallel_attributes,
+)
+
+from apex.transformer.tensor_parallel.mappings import (
+    copy_to_tensor_model_parallel_region,
+    gather_from_tensor_model_parallel_region,
+    reduce_from_tensor_model_parallel_region,
+    scatter_to_tensor_model_parallel_region,
+)
+
+from .random import (
+    checkpoint,
+    get_cuda_rng_tracker,
+    init_checkpointed_activations_memory_buffer,
+    model_parallel_cuda_manual_seed,
+    reset_checkpointed_activations_memory_buffer,
+)
+
+from apex.transformer.tensor_parallel.utils import split_tensor_along_last_dim
+
+
+__all__ = [
+    # cross_entropy.py
+    "vocab_parallel_cross_entropy",
+    # data.py
+    "broadcast_data",
+    # layers.py
+    "ColumnParallelLinear",
+    "RowParallelLinear",
+    "VocabParallelEmbedding",
+    "set_tensor_model_parallel_attributes",
+    "set_defaults_if_not_set_tensor_model_parallel_attributes",
+    "copy_tensor_model_parallel_attributes",
+    # mappings.py
+    "copy_to_tensor_model_parallel_region",
+    "gather_from_tensor_model_parallel_region",
+    "reduce_from_tensor_model_parallel_region",
+    "scatter_to_tensor_model_parallel_region",
+    # random.py
+    "checkpoint",
+    "get_cuda_rng_tracker",
+    "init_checkpointed_activations_memory_buffer",
+    "model_parallel_cuda_manual_seed",
+    "reset_checkpointed_activations_memory_buffer",
+    # utils.py
+    "split_tensor_along_last_dim",
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/cross_entropy.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/cross_entropy.py
new file mode 100644
index 0000000000..3918645c71
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/cross_entropy.py
@@ -0,0 +1,103 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import torch
+
+from apex.transformer.parallel_state import get_tensor_model_parallel_group
+from apex.transformer.parallel_state import get_tensor_model_parallel_rank
+from apex.transformer.parallel_state import get_tensor_model_parallel_world_size
+from apex.transformer.tensor_parallel.utils import VocabUtility
+
+
+class _VocabParallelCrossEntropy(torch.autograd.Function):
+    @staticmethod
+    def forward(ctx, vocab_parallel_logits, target):
+
+        # Maximum value along vocab dimension across all GPUs.
+        logits_max = torch.max(vocab_parallel_logits, dim=-1)[0]
+        torch.distributed.all_reduce(
+            logits_max, op=torch.distributed.ReduceOp.MAX, group=get_tensor_model_parallel_group()
+        )
+        # Subtract the maximum value.
+        vocab_parallel_logits = vocab_parallel_logits - logits_max.unsqueeze(dim=-1)
+
+        # Get the partition's vocab indecies
+        get_vocab_range = VocabUtility.vocab_range_from_per_partition_vocab_size
+        partition_vocab_size = vocab_parallel_logits.size()[-1]
+        rank = get_tensor_model_parallel_rank()
+        world_size = get_tensor_model_parallel_world_size()
+        vocab_start_index, vocab_end_index = get_vocab_range(partition_vocab_size, rank, world_size)
+
+        # Create a mask of valid vocab ids (1 means it needs to be masked).
+        target_mask = (target < vocab_start_index) | (target >= vocab_end_index)
+        masked_target = target.clone() - vocab_start_index
+        masked_target[target_mask] = 0
+
+        # Get predicted-logits = logits[target].
+        # For Simplicity, we convert logits to a 2-D tensor with size
+        # [*, partition-vocab-size] and target to a 1-D tensor of size [*].
+        logits_2d = vocab_parallel_logits.view(-1, partition_vocab_size)
+        masked_target_1d = masked_target.view(-1)
+        arange_1d = torch.arange(start=0, end=logits_2d.size()[0], device=logits_2d.device)
+        predicted_logits_1d = logits_2d[arange_1d, masked_target_1d]
+        predicted_logits_1d = predicted_logits_1d.clone().contiguous()
+        predicted_logits = predicted_logits_1d.view_as(target)
+        predicted_logits[target_mask] = 0.0
+        # All reduce is needed to get the chunks from other GPUs.
+        torch.distributed.all_reduce(
+            predicted_logits, op=torch.distributed.ReduceOp.SUM, group=get_tensor_model_parallel_group()
+        )
+
+        # Sum of exponential of logits along vocab dimension across all GPUs.
+        exp_logits = vocab_parallel_logits
+        torch.exp(vocab_parallel_logits, out=exp_logits)
+        sum_exp_logits = exp_logits.sum(dim=-1)
+        torch.distributed.all_reduce(
+            sum_exp_logits, op=torch.distributed.ReduceOp.SUM, group=get_tensor_model_parallel_group()
+        )
+
+        # Loss = log(sum(exp(logits))) - predicted-logit.
+        loss = torch.log(sum_exp_logits) - predicted_logits
+
+        # Store softmax, target-mask and masked-target for backward pass.
+        exp_logits.div_(sum_exp_logits.unsqueeze(dim=-1))
+        ctx.save_for_backward(exp_logits, target_mask, masked_target_1d)
+
+        return loss
+
+    @staticmethod
+    def backward(ctx, grad_output):
+
+        # Retreive tensors from the forward path.
+        softmax, target_mask, masked_target_1d = ctx.saved_tensors
+
+        # All the inputs have softmax as thier gradient.
+        grad_input = softmax
+        # For simplicity, work with the 2D gradient.
+        partition_vocab_size = softmax.size()[-1]
+        grad_2d = grad_input.view(-1, partition_vocab_size)
+
+        # Add the gradient from matching classes.
+        arange_1d = torch.arange(start=0, end=grad_2d.size()[0], device=grad_2d.device)
+        grad_2d[arange_1d, masked_target_1d] -= 1.0 - target_mask.view(-1).float()
+
+        # Finally elementwise multiplication with the output gradients.
+        grad_input.mul_(grad_output.unsqueeze(dim=-1))
+
+        return grad_input, None
+
+
+def vocab_parallel_cross_entropy(vocab_parallel_logits, target):
+    """Helper function for the cross entropy."""
+    return _VocabParallelCrossEntropy.apply(vocab_parallel_logits, target)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/data.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/data.py
new file mode 100644
index 0000000000..558d289413
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/data.py
@@ -0,0 +1,113 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import torch
+
+from apex.transformer.parallel_state import get_tensor_model_parallel_group
+from apex.transformer.parallel_state import get_tensor_model_parallel_rank
+from apex.transformer.parallel_state import get_tensor_model_parallel_src_rank
+
+
+_MAX_DATA_DIM = 5
+
+
+def _check_data_types(keys, data, target_dtype):
+    """Check that all the keys have the same target data type."""
+    for key in keys:
+        assert data[key].dtype == target_dtype, "{} has data type {} which " "is different than {}".format(
+            key, data[key].dtype, target_dtype
+        )
+
+
+def _build_key_size_numel_dictionaries(keys, data):
+    """Build the size on rank 0 and broadcast."""
+    max_dim = _MAX_DATA_DIM
+    sizes = [0 for _ in range(max_dim) for _ in keys]
+
+    # Pack the sizes on rank zero.
+    if get_tensor_model_parallel_rank() == 0:
+        offset = 0
+        for key in keys:
+            assert data[key].dim() < max_dim, "you should increase MAX_DATA_DIM"
+            size = data[key].size()
+            for i, s in enumerate(size):
+                sizes[i + offset] = s
+            offset += max_dim
+
+    # Move to GPU and broadcast.
+    sizes_cuda = torch.cuda.LongTensor(sizes)
+    torch.distributed.broadcast(
+        sizes_cuda, get_tensor_model_parallel_src_rank(), group=get_tensor_model_parallel_group(),
+    )
+
+    # Move back to cpu and unpack.
+    sizes_cpu = sizes_cuda.cpu()
+    key_size = {}
+    key_numel = {}
+    total_numel = 0
+    offset = 0
+    for key in keys:
+        i = 0
+        size = []
+        numel = 1
+        while sizes_cpu[offset + i] > 0:
+            this_size = sizes_cpu[offset + i]
+            size.append(this_size)
+            numel *= this_size
+            i += 1
+        key_size[key] = size
+        key_numel[key] = numel
+        total_numel += numel
+        offset += max_dim
+
+    return key_size, key_numel, total_numel
+
+
+def broadcast_data(keys, data, datatype):
+    """Broadcast data from rank zero of each model parallel group to the
+    members of the same model parallel group.
+
+    Arguments:
+        keys: list of keys in the data disctionary to be broadcasted
+        data: data dictionary of string keys and cpu tensor values.
+        datatype: torch data type of all tensors in data associated
+                  with keys.
+    """
+    # Build (key, size) and (key, number of elements) dictionaries along
+    # with the total number of elements on all ranks.
+    key_size, key_numel, total_numel = _build_key_size_numel_dictionaries(keys, data)
+    # Pack on rank zero.
+    if get_tensor_model_parallel_rank() == 0:
+        # Check that all keys have the same data type.
+        _check_data_types(keys, data, datatype)
+        # Flatten the data associated with the keys
+        flatten_data = torch.cat([data[key].contiguous().view(-1) for key in keys], dim=0).cuda()
+    else:
+        flatten_data = torch.empty(total_numel, device=torch.cuda.current_device(), dtype=datatype)
+
+    # Broadcast
+    torch.distributed.broadcast(
+        flatten_data, get_tensor_model_parallel_src_rank(), group=get_tensor_model_parallel_group(),
+    )
+
+    # Unpack
+    output = {}
+    offset = 0
+    for key in keys:
+        size = key_size[key]
+        numel = key_numel[key]
+        output[key] = flatten_data.narrow(0, offset, numel).view(size)
+        offset += numel
+
+    return output
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/layers.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/layers.py
new file mode 100644
index 0000000000..9264d63777
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/layers.py
@@ -0,0 +1,477 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+
+# Parts of the code here are adapted from PyTorch
+# repo: https://github.com/pytorch/pytorch
+import torch
+import torch.nn.functional as F
+import torch.nn.init as init
+from torch.nn.parameter import Parameter
+
+from apex._autocast_utils import _cast_if_autocast_enabled
+from apex.transformer.parallel_state import get_tensor_model_parallel_group
+from apex.transformer.parallel_state import get_tensor_model_parallel_rank
+from apex.transformer.parallel_state import get_tensor_model_parallel_world_size
+from apex.transformer.utils import divide
+from apex.transformer.tensor_parallel.mappings import copy_to_tensor_model_parallel_region
+from apex.transformer.tensor_parallel.mappings import gather_from_tensor_model_parallel_region
+from apex.transformer.tensor_parallel.mappings import reduce_from_tensor_model_parallel_region
+from apex.transformer.tensor_parallel.mappings import scatter_to_tensor_model_parallel_region
+from apex.transformer.tensor_parallel.random import get_cuda_rng_tracker
+from apex.transformer.tensor_parallel.utils import VocabUtility
+
+
+_MODEL_PARALLEL_ATTRIBUTE_DEFAULTS = {
+    "tensor_model_parallel": False,
+    "partition_dim": -1,
+    "partition_stride": 1,
+}
+
+
+def param_is_not_tensor_parallel_duplicate(param):
+    return (hasattr(param, "tensor_model_parallel") and param.tensor_model_parallel) or (
+        get_tensor_model_parallel_rank() == 0
+    )
+
+
+def set_tensor_model_parallel_attributes(tensor, is_parallel, dim, stride):
+    # Make sure the attributes are not set.
+    for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS:
+        assert not hasattr(tensor, attribute)
+    # Set the attributes.
+    setattr(tensor, "tensor_model_parallel", is_parallel)
+    setattr(tensor, "partition_dim", dim)
+    setattr(tensor, "partition_stride", stride)
+
+
+def set_defaults_if_not_set_tensor_model_parallel_attributes(tensor):
+    def maybe_set(attribute, value):
+        if not hasattr(tensor, attribute):
+            setattr(tensor, attribute, value)
+
+    for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS:
+        maybe_set(attribute, _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS[attribute])
+
+
+def copy_tensor_model_parallel_attributes(destination_tensor, source_tensor):
+    def maybe_copy(attribute):
+        if hasattr(source_tensor, attribute):
+            setattr(destination_tensor, attribute, getattr(source_tensor, attribute))
+
+    for attribute in _MODEL_PARALLEL_ATTRIBUTE_DEFAULTS:
+        maybe_copy(attribute)
+
+
+def _initialize_affine_weight_gpu(weight, init_method, partition_dim, stride=1):
+    """Initialize affine weight for model parallel on GPU."""
+
+    set_tensor_model_parallel_attributes(tensor=weight, is_parallel=True, dim=partition_dim, stride=stride)
+
+    with get_cuda_rng_tracker().fork():
+        init_method(weight)
+
+
+# TODO (mkozuki): Re-consider removing params_dtype from arguments to make this
+# more parallel with _initialize_affine_weight_gpu
+def _initialize_affine_weight_cpu(
+    weight,
+    output_size,
+    input_size,
+    per_partition_size,
+    partition_dim,
+    init_method,
+    stride=1,
+    return_master_weight=False,
+    *,
+    params_dtype=torch.float32,
+):
+    """Initialize affine weight for model parallel.
+
+    Build the master weight on all processes and scatter
+    the relevant chunk."""
+
+    set_tensor_model_parallel_attributes(tensor=weight, is_parallel=True, dim=partition_dim, stride=stride)
+
+    # Initialize master weight
+    master_weight = torch.empty(output_size, input_size, dtype=torch.float, requires_grad=False)
+    init_method(master_weight)
+    master_weight = master_weight.to(dtype=params_dtype)
+
+    # Split and copy
+    per_partition_per_stride_size = divide(per_partition_size, stride)
+    weight_list = torch.split(master_weight, per_partition_per_stride_size, dim=partition_dim)
+    rank = get_tensor_model_parallel_rank()
+    world_size = get_tensor_model_parallel_world_size()
+    my_weight_list = weight_list[rank::world_size]
+
+    with torch.no_grad():
+        torch.cat(my_weight_list, dim=partition_dim, out=weight)
+    if return_master_weight:
+        return master_weight
+    return None
+
+
+class VocabParallelEmbedding(torch.nn.Module):
+    """Embedding parallelized in the vocabulary dimension.
+
+    This is mainly adapted from torch.nn.Embedding and all the default
+    values are kept.
+    Arguments:
+        num_embeddings: vocabulary size.
+        embedding_dim: size of hidden state.
+        init_method: method to initialize weights.
+    """
+
+    def __init__(
+        self, num_embeddings, embedding_dim, init_method=init.xavier_normal_, *, params_dtype=torch.float32, use_cpu_initialization=False,
+    ):
+        super(VocabParallelEmbedding, self).__init__()
+        # Keep the input dimensions.
+        self.num_embeddings = num_embeddings
+        self.embedding_dim = embedding_dim
+        # Set the detauls for compatibility.
+        self.padding_idx = None
+        self.max_norm = None
+        self.norm_type = 2.0
+        self.scale_grad_by_freq = False
+        self.sparse = False
+        self._weight = None
+        self.tensor_model_parallel_size = get_tensor_model_parallel_world_size()
+        # Divide the weight matrix along the vocaburaly dimension.
+        self.vocab_start_index, self.vocab_end_index = VocabUtility.vocab_range_from_global_vocab_size(
+            self.num_embeddings, get_tensor_model_parallel_rank(), self.tensor_model_parallel_size
+        )
+        self.num_embeddings_per_partition = self.vocab_end_index - self.vocab_start_index
+
+        # Allocate weights and initialize.
+        if use_cpu_initialization:
+            self.weight = Parameter(
+                torch.empty(self.num_embeddings_per_partition, self.embedding_dim, dtype=params_dtype)
+            )
+            _initialize_affine_weight_cpu(
+                self.weight, self.num_embeddings, self.embedding_dim, self.num_embeddings_per_partition, 0, init_method,
+                params_dtype=params_dtype,
+            )
+        else:
+            self.weight = Parameter(
+                torch.empty(
+                    self.num_embeddings_per_partition,
+                    self.embedding_dim,
+                    device=torch.cuda.current_device(),
+                    dtype=params_dtype,
+                )
+            )
+            _initialize_affine_weight_gpu(self.weight, init_method, partition_dim=0, stride=1)
+
+    def forward(self, input_):
+        if self.tensor_model_parallel_size > 1:
+            # Build the mask.
+            input_mask = (input_ < self.vocab_start_index) | (input_ >= self.vocab_end_index)
+            # Mask the input.
+            masked_input = input_.clone() - self.vocab_start_index
+            masked_input[input_mask] = 0
+        else:
+            masked_input = input_
+            # Get the embeddings.
+        output_parallel = F.embedding(
+            masked_input,
+            self.weight,
+            self.padding_idx,
+            self.max_norm,
+            self.norm_type,
+            self.scale_grad_by_freq,
+            self.sparse,
+        )
+        # Mask the output embedding.
+        if self.tensor_model_parallel_size > 1:
+            output_parallel[input_mask, :] = 0.0
+        # Reduce across all the model parallel GPUs.
+        output = reduce_from_tensor_model_parallel_region(output_parallel)
+        return output
+
+
+class ColumnParallelLinearWithAsyncAllreduce(torch.autograd.Function):
+    """
+    Column-parallel linear layer execution with asynchronous all-reduce
+    execution in backprop.
+    """
+    @staticmethod
+    def forward(ctx, input, weight, bias):
+        ctx.save_for_backward(input, weight)
+        ctx.use_bias = bias is not None
+        output = torch.matmul(input, weight.t())
+        if bias is not None:
+            output = output + bias
+        return output
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, weight = ctx.saved_tensors
+        use_bias = ctx.use_bias
+        grad_input = grad_output.matmul(weight)
+        # Asynchronous all-reduce
+        handle = torch.distributed.all_reduce(
+                grad_input, group=get_tensor_model_parallel_group(), async_op=True)
+        # Delay the start of weight gradient computation shortly (3us) to have
+        # all-reduce scheduled first and have GPU resources allocated
+        _ = torch.empty(1, device=grad_output.device) + 1
+        grad_weight = grad_output.t().matmul(input)
+        grad_bias = grad_output.sum(dim=0) if use_bias else None
+        handle.wait()
+        return grad_input, grad_weight, grad_bias
+
+
+def column_parallel_linear(input, weight, bias):
+    args = _cast_if_autocast_enabled(input, weight, bias)
+    with torch.cuda.amp.autocast(enabled=False):
+        return ColumnParallelLinearWithAsyncAllreduce.apply(*args)
+
+
+class ColumnParallelLinear(torch.nn.Module):
+    """Linear layer with column parallelism.
+
+    The linear layer is defined as Y = XA + b. A is parallelized along
+    its second dimension as A = [A_1, ..., A_p].
+
+    Arguments:
+        input_size: first dimension of matrix A.
+        output_size: second dimension of matrix A.
+        bias: If true, add bias
+        gather_output: If true, call all-gether on output and make Y avaiable
+                       to all GPUs, otherwise, every GPU will have its output
+                       which is Y_i = XA_i
+        init_method: method to initialize weights. Note that bias is always set
+                     to zero.
+        stride: For the strided linear layers.
+        keep_master_weight_for_test: This was added for testing and should be
+                                     set to False. It returns the master weights
+                                     used for initialization.
+        skip_bias_add: This was added to enable performance optimations where bias
+                       can be fused with other elementwise operations. we skip
+                       adding bias but instead return it.
+    """
+
+    def __init__(
+        self,
+        input_size,
+        output_size,
+        bias=True,
+        gather_output=True,
+        init_method=init.xavier_normal_,
+        stride=1,
+        keep_master_weight_for_test=False,
+        skip_bias_add=False,
+        *,
+        no_async_tensor_model_parallel_allreduce=False,
+        params_dtype=torch.float32,
+        use_cpu_initialization=False,
+    ):
+        super(ColumnParallelLinear, self).__init__()
+
+        # Keep input parameters
+        self.input_size = input_size
+        self.output_size = output_size
+        self.gather_output = gather_output
+        # Divide the weight matrix along the last dimension.
+        world_size = get_tensor_model_parallel_world_size()
+        self.output_size_per_partition = divide(output_size, world_size)
+        self.skip_bias_add = skip_bias_add
+
+        # Parameters.
+        # Note: torch.nn.functional.linear performs XA^T + b and as a result
+        # we allocate the transpose.
+        # Initialize weight.
+        if use_cpu_initialization:
+            self.weight = Parameter(torch.empty(self.output_size_per_partition, self.input_size, dtype=params_dtype))
+            self.master_weight = _initialize_affine_weight_cpu(
+                self.weight,
+                self.output_size,
+                self.input_size,
+                self.output_size_per_partition,
+                0,
+                init_method,
+                stride=stride,
+                return_master_weight=keep_master_weight_for_test,
+                params_dtype=params_dtype,
+            )
+        else:
+            self.weight = Parameter(
+                torch.empty(
+                    self.output_size_per_partition,
+                    self.input_size,
+                    device=torch.cuda.current_device(),
+                    dtype=params_dtype,
+                )
+            )
+            _initialize_affine_weight_gpu(self.weight, init_method, partition_dim=0, stride=stride)
+
+        if bias:
+            if use_cpu_initialization:
+                self.bias = Parameter(torch.empty(self.output_size_per_partition, dtype=params_dtype))
+            else:
+                self.bias = Parameter(
+                    torch.empty(self.output_size_per_partition, device=torch.cuda.current_device(), dtype=params_dtype)
+                )
+            set_tensor_model_parallel_attributes(self.bias, True, 0, stride)
+            # Always initialize bias to zero.
+            with torch.no_grad():
+                self.bias.zero_()
+        else:
+            self.register_parameter("bias", None)
+
+        self.async_tensor_model_parallel_allreduce = (
+                not no_async_tensor_model_parallel_allreduce and
+                world_size > 1)
+
+    def forward(self, input_):
+        bias = self.bias if not self.skip_bias_add else None
+
+        if self.async_tensor_model_parallel_allreduce:
+            input_shape = input_.shape
+            input_ = input_.view(input_shape[0] * input_shape[1],input_shape[2])
+            # Matrix multiply with asynchronous all-reduce execution
+            output_parallel = column_parallel_linear(input_, self.weight, bias)
+            output_parallel = output_parallel.view(
+                    input_shape[0], input_shape[1], output_parallel.shape[1])
+        else:
+            # Set up backprop all-reduce.
+            input_parallel = copy_to_tensor_model_parallel_region(input_)
+
+            # Matrix multiply.
+            output_parallel = F.linear(input_parallel, self.weight, bias)
+
+        if self.gather_output:
+            # All-gather across the partitions.
+            output = gather_from_tensor_model_parallel_region(output_parallel)
+        else:
+            output = output_parallel
+        output_bias = self.bias if self.skip_bias_add else None
+        return output, output_bias
+
+
+class RowParallelLinear(torch.nn.Module):
+    """Linear layer with row parallelism.
+
+    The linear layer is defined as Y = XA + b. A is parallelized along
+    its first dimension and X along its second dimension as:
+               -   -
+              | A_1 |
+              | .   |
+          A = | .   |        X = [X_1, ..., X_p]
+              | .   |
+              | A_p |
+               -   -
+    Arguments:
+        input_size: first dimension of matrix A.
+        output_size: second dimension of matrix A.
+        bias: If true, add bias. Note that bias is not parallelized.
+        input_is_parallel: If true, we assume that the input is already
+                           split across the GPUs and we do not split
+                           again.
+        init_method: method to initialize weights. Note that bias is always set
+                     to zero.
+        stride: For the strided linear layers.
+        keep_master_weight_for_test: This was added for testing and should be
+                                     set to False. It returns the master weights
+                                     used for initialization.
+        skip_bias_add: This was added to enable performance optimization where bias
+                       can be fused with other elementwise operations. We skip
+                       adding bias but instead return it.
+    """
+
+    def __init__(
+        self,
+        input_size,
+        output_size,
+        bias=True,
+        input_is_parallel=False,
+        init_method=init.xavier_normal_,
+        stride=1,
+        keep_master_weight_for_test=False,
+        skip_bias_add=False,
+        *,
+        params_dtype=torch.float32,
+        use_cpu_initialization=False,
+    ):
+        super(RowParallelLinear, self).__init__()
+
+        # Keep input parameters
+        self.input_size = input_size
+        self.output_size = output_size
+        self.input_is_parallel = input_is_parallel
+        # Divide the weight matrix along the last dimension.
+        world_size = get_tensor_model_parallel_world_size()
+        self.input_size_per_partition = divide(input_size, world_size)
+        self.skip_bias_add = skip_bias_add
+
+        # as an argument to this function?
+        # Parameters.
+        # Note: torch.nn.functional.linear performs XA^T + b and as a result
+        # we allocate the transpose.
+        # Initialize weight.
+        if use_cpu_initialization:
+            self.weight = Parameter(torch.empty(self.output_size, self.input_size_per_partition, dtype=params_dtype))
+            self.master_weight = _initialize_affine_weight_cpu(
+                self.weight,
+                self.output_size,
+                self.input_size,
+                self.input_size_per_partition,
+                1,
+                init_method,
+                stride=stride,
+                return_master_weight=keep_master_weight_for_test,
+                params_dtype=params_dtype,
+            )
+        else:
+            self.weight = Parameter(
+                torch.empty(
+                    self.output_size,
+                    self.input_size_per_partition,
+                    device=torch.cuda.current_device(),
+                    dtype=params_dtype,
+                )
+            )
+            _initialize_affine_weight_gpu(self.weight, init_method, partition_dim=1, stride=stride)
+        if bias:
+            if use_cpu_initialization:
+                self.bias = Parameter(torch.empty(self.output_size, dtype=params_dtype))
+            else:
+                self.bias = Parameter(
+                    torch.empty(self.output_size, device=torch.cuda.current_device(), dtype=params_dtype)
+                )
+            # Always initialize bias to zero.
+            with torch.no_grad():
+                self.bias.zero_()
+        else:
+            self.register_parameter("bias", None)
+
+    def forward(self, input_):
+        # Set up backprop all-reduce.
+        if self.input_is_parallel:
+            input_parallel = input_
+        else:
+            input_parallel = scatter_to_tensor_model_parallel_region(input_)
+        # Matrix multiply.
+        output_parallel = F.linear(input_parallel, self.weight)
+        # All-reduce across all the partitions.
+        output_ = reduce_from_tensor_model_parallel_region(output_parallel)
+        if not self.skip_bias_add:
+            output = output_ + self.bias if self.bias is not None else output_
+            output_bias = None
+        else:
+            output = output_
+            output_bias = self.bias
+        return output, output_bias
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/mappings.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/mappings.py
new file mode 100644
index 0000000000..882f72de1f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/mappings.py
@@ -0,0 +1,159 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import torch
+
+from apex.transformer.parallel_state import get_tensor_model_parallel_group
+from apex.transformer.parallel_state import get_tensor_model_parallel_world_size
+from apex.transformer.parallel_state import get_tensor_model_parallel_rank
+from apex.transformer.tensor_parallel.utils import split_tensor_along_last_dim
+
+
+def _reduce(input_):
+    """All-reduce the input tensor across model parallel group."""
+
+    # Bypass the function if we are using only 1 GPU.
+    if get_tensor_model_parallel_world_size() == 1:
+        return input_
+
+    # All-reduce.
+    torch.distributed.all_reduce(input_, group=get_tensor_model_parallel_group())
+
+    return input_
+
+
+def _split(input_):
+    """Split the tensor along its last dimension and keep the
+    corresponding slice."""
+
+    world_size = get_tensor_model_parallel_world_size()
+    # Bypass the function if we are using only 1 GPU.
+    if world_size == 1:
+        return input_
+
+    # Split along last dimension.
+    input_list = split_tensor_along_last_dim(input_, world_size)
+
+    # Note: torch.split does not create contiguous tensors by default.
+    rank = get_tensor_model_parallel_rank()
+    output = input_list[rank].contiguous()
+
+    return output
+
+
+def _gather(input_):
+    """Gather tensors and concatinate along the last dimension."""
+
+    world_size = get_tensor_model_parallel_world_size()
+    # Bypass the function if we are using only 1 GPU.
+    if world_size == 1:
+        return input_
+
+    # Size and dimension.
+    last_dim = input_.dim() - 1
+    rank = get_tensor_model_parallel_rank()
+
+    tensor_list = [torch.empty_like(input_) for _ in range(world_size)]
+    tensor_list[rank] = input_
+    torch.distributed.all_gather(tensor_list, input_, group=get_tensor_model_parallel_group())
+
+    # Note: torch.cat already creates a contiguous tensor.
+    output = torch.cat(tensor_list, dim=last_dim).contiguous()
+
+    return output
+
+
+class _CopyToModelParallelRegion(torch.autograd.Function):
+    """Pass the input to the model parallel region."""
+
+    @staticmethod
+    def symbolic(graph, input_):
+        return input_
+
+    @staticmethod
+    def forward(ctx, input_):
+        return input_
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return _reduce(grad_output)
+
+
+class _ReduceFromModelParallelRegion(torch.autograd.Function):
+    """All-reduce the input from the model parallel region."""
+
+    @staticmethod
+    def symbolic(graph, input_):
+        return _reduce(input_)
+
+    @staticmethod
+    def forward(ctx, input_):
+        return _reduce(input_)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return grad_output
+
+
+class _ScatterToModelParallelRegion(torch.autograd.Function):
+    """Split the input and keep only the corresponding chuck to the rank."""
+
+    @staticmethod
+    def symbolic(graph, input_):
+        return _split(input_)
+
+    @staticmethod
+    def forward(ctx, input_):
+        return _split(input_)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return _gather(grad_output)
+
+
+class _GatherFromModelParallelRegion(torch.autograd.Function):
+    """Gather the input from model parallel region and concatinate."""
+
+    @staticmethod
+    def symbolic(graph, input_):
+        return _gather(input_)
+
+    @staticmethod
+    def forward(ctx, input_):
+        return _gather(input_)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        return _split(grad_output)
+
+
+# -----------------
+# Helper functions.
+# -----------------
+
+
+def copy_to_tensor_model_parallel_region(input_):
+    return _CopyToModelParallelRegion.apply(input_)
+
+
+def reduce_from_tensor_model_parallel_region(input_):
+    return _ReduceFromModelParallelRegion.apply(input_)
+
+
+def scatter_to_tensor_model_parallel_region(input_):
+    return _ScatterToModelParallelRegion.apply(input_)
+
+
+def gather_from_tensor_model_parallel_region(input_):
+    return _GatherFromModelParallelRegion.apply(input_)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/memory.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/memory.py
new file mode 100644
index 0000000000..51d526488f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/memory.py
@@ -0,0 +1,136 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import torch
+
+
+# A dictionary of all the memory buffers allocated.
+_MEM_BUFFS = dict()
+
+
+def allocate_mem_buff(name, numel, dtype, track_usage):
+    """Allocate a memory buffer."""
+    assert name not in _MEM_BUFFS, "memory buffer {} already allocated.".format(name)
+    _MEM_BUFFS[name] = MemoryBuffer(name, numel, dtype, track_usage)
+    return _MEM_BUFFS[name]
+
+
+def get_mem_buff(name):
+    """Get the memory buffer."""
+    return _MEM_BUFFS[name]
+
+
+class MemoryBuffer:
+    """Contiguous memory buffer.
+    Allocate a contiguous memory of type `dtype` and size `numel`. It is
+    used to reduce memory fragmentation.
+
+    Usage: After the allocation, the `_start` index is set tot the first
+           index of the memory. A memory chunk starting from `_start` index
+           can be `allocated` for an input tensor, with the elements of the
+           tensor being coppied. The buffer can be reused by resetting the
+           `_start` index.
+
+    """
+
+    def __init__(self, name, numel, dtype, track_usage):
+        if torch.distributed.get_rank() == 0:
+            element_size = torch.tensor([], dtype=dtype).element_size()
+            print(
+                "> building the {} memory buffer with {} num elements "
+                "and {} dtype ({:.1f} MB)...".format(name, numel, dtype, numel * element_size / 1024 / 1024),
+                flush=True,
+            )
+        self.name = name
+        self.numel = numel
+        self.dtype = dtype
+        self.data = torch.empty(self.numel, dtype=self.dtype, device=torch.cuda.current_device(), requires_grad=False)
+
+        # Index tracking the start of the free memory.
+        self._start = 0
+
+        # Values used for tracking usage.
+        self.track_usage = track_usage
+        if self.track_usage:
+            self.in_use_value = 0.0
+            self.total_value = 0.0
+
+    def reset(self):
+        """Reset the buffer start index to the beginning of the buffer."""
+        self._start = 0
+
+    def is_in_use(self):
+        """Whether the current buffer hold on to any memory."""
+        return self._start > 0
+
+    def numel_in_use(self):
+        """Return number of elements in use."""
+        return self._start
+
+    def add(self, tensor):
+        """Allocate a chunk of memory from the buffer to tensor and copy
+        the values."""
+        assert tensor.dtype == self.dtype, "Input tensor type {} different from buffer type {}".format(
+            tensor.dtype, self.dtype
+        )
+        # Number of elements of the input tensor.
+        tensor_numel = torch.numel(tensor)
+        new_start = self._start + tensor_numel
+        assert new_start <= self.numel, "Not enough memory left in the buffer ({} > {})".format(
+            tensor_numel, self.numel - self._start
+        )
+        # New tensor is a view into the memory.
+        new_tensor = self.data[self._start : new_start]
+        self._start = new_start
+        new_tensor = new_tensor.view(tensor.shape)
+        new_tensor.copy_(tensor)
+        # Return a pointer to the new tensor.
+        return new_tensor
+
+    def get_data(self):
+        """Return the data currently in use."""
+        if self.track_usage:
+            self.in_use_value += float(self._start)
+            self.total_value += float(self.numel)
+        return self.data[: self._start]
+
+    def print_average_usage(self):
+        """Print memory usage average over time. We would like this value
+        to be as high as possible."""
+        assert self.track_usage, "You need to enable track usage."
+        if torch.distributed.get_rank() == 0:
+            print(
+                " > usage of {} memory buffer: {:.2f} %".format(
+                    self.name, self.in_use_value * 100.0 / self.total_value
+                ),
+                flush=True,
+            )
+
+
+class RingMemBuffer:
+    """A ring of memory buffers."""
+
+    def __init__(self, name, num_buffers, numel, dtype, track_usage):
+        self.num_buffers = num_buffers
+        self.buffers = [
+            allocate_mem_buff(name + " {}".format(i), numel, dtype, track_usage) for i in range(num_buffers)
+        ]
+        self._index = -1
+
+    def get_next_buffer(self):
+        self._index += 1
+        self._index = self._index % self.num_buffers
+        buff = self.buffers[self._index]
+        assert not buff.is_in_use(), "buffer is already in use."
+        return buff
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/random.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/random.py
new file mode 100644
index 0000000000..c94d8d25c0
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/random.py
@@ -0,0 +1,294 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+# TODO (mkozuki): Audit this file.
+# I don't think some functions strongly relate to `random` in tensor_parallel.
+# Rather, some functions are mainly for gradient checkpointing (torch.utils.checkpoint).
+
+# Parts of the code here are adapted from PyTorch
+# repo: https://github.com/pytorch/pytorch
+import contextlib
+
+import torch
+from torch import _C
+from torch.cuda import _lazy_call, device as device_ctx_manager
+from torch.utils.checkpoint import detach_variable
+
+from apex.transformer.parallel_state import get_tensor_model_parallel_rank
+from apex.transformer.tensor_parallel.memory import allocate_mem_buff
+from apex.transformer.utils import split_tensor_into_1d_equal_chunks
+from apex.transformer.utils import gather_split_1d_tensor
+
+
+# Default name for the model parallel rng tracker.
+_MODEL_PARALLEL_RNG_TRACKER_NAME = "model-parallel-rng"
+
+
+# Whether apply model parallelism to checkpointed hidden states.
+_CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER = None
+
+
+# TODO (mkozuki): Consider the possibility of removing `tensor_model_parallel_size`,
+# `get_tensor_model_parallel_world_size()` might be alternative.
+def init_checkpointed_activations_memory_buffer(
+    micro_batch_size,
+    max_position_embeddings,
+    hidden_size,
+    num_layers,
+    tensor_model_parallel_size,
+    checkpoint_num_layers,
+    fp16,
+):
+    """Initializ the memory buffer for the checkpointed activations."""
+
+    per_layer = micro_batch_size * max_position_embeddings * hidden_size // tensor_model_parallel_size
+    assert num_layers % checkpoint_num_layers == 0, "number of layers is not divisible by checkpoint-num-layers"
+    num_checkpointer_layers = num_layers // checkpoint_num_layers
+    numel = per_layer * num_checkpointer_layers
+    dtype = torch.half
+    if not fp16:
+        dtype = torch.float
+
+    global _CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER
+    assert (
+        _CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER is None
+    ), "checkpointed activations memory buffer is already allocated."
+    _CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER = allocate_mem_buff(
+        "checkpointed activations", numel, dtype, track_usage=False
+    )
+
+
+def reset_checkpointed_activations_memory_buffer():
+    """Reset the memory used for checkpointing."""
+    if _CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER is not None:
+        _CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER.reset()
+
+
+def _set_cuda_rng_state(new_state, device=-1):
+    """Sets the random number generator state of the current GPU.
+
+    Argumentss:
+        new_state (torch.ByteTensor): The desired state
+    This function is adapted from PyTorch repo (torch.cuda.set_rng_state)
+    with a single change: the input state is not cloned. Cloning caused
+    major performance issues for +4 GPU cases.
+    """
+    if hasattr(_C, "_cuda_setRNGState") and callable(_C._cuda_setRNGState):
+        # older PyTorch
+        def cb():
+            with device_ctx_manager(device):
+                _C._cuda_setRNGState(new_state)
+
+    else:
+        # newer PyTorch
+        if device == -1:
+            device = torch.device("cuda")
+        elif isinstance(device, str):
+            device = torch.device(device)
+        elif isinstance(device, int):
+            device = torch.device("cuda", device)
+
+        def cb():
+            idx = device.index
+            if idx is None:
+                idx = torch.cuda.current_device()
+            default_generator = torch.cuda.default_generators[idx]
+            default_generator.set_state(new_state)
+
+    _lazy_call(cb)
+
+
+class CudaRNGStatesTracker:
+    """Tracker for the cuda RNG states.
+
+    Using the `add` method, a cuda rng state is initialized based on
+    the input `seed` and is assigned to `name`. Later, by forking the
+    rng state, we can perform operations and return to our starting
+    cuda state.
+    """
+
+    def __init__(self):
+        # Map from a string name to the cuda rng state.
+        self.states_ = {}
+        # Seeds are just for book keeping and ensure no seed is set twice.
+        self.seeds_ = set()
+
+    def reset(self):
+        """Set to the initial state (no tracker)."""
+        self.states_ = {}
+        self.seeds_ = set()
+
+    def get_states(self):
+        """Get rng states. Copy the dictionary so we have direct
+        pointers to the states, not just a pointer to the dictionary."""
+        states = {}
+        for name in self.states_:
+            states[name] = self.states_[name]
+        return states
+
+    def set_states(self, states):
+        """Set the rng states. For efficiency purposes, we do not check
+        the size of seed for compatibility."""
+        self.states_ = states
+
+    def add(self, name, seed):
+        """Track the rng state."""
+        # Check seed is not already used.
+        if seed in self.seeds_:
+            raise Exception("seed {} already exists".format(seed))
+        self.seeds_.add(seed)
+        # Check that state is not already defined.
+        if name in self.states_:
+            raise Exception("cuda rng state {} already exists".format(name))
+        # Get the current rng state.
+        orig_rng_state = torch.cuda.get_rng_state()
+        # Set the new state and store it.
+        torch.cuda.manual_seed(seed)
+        self.states_[name] = torch.cuda.get_rng_state()
+        # Reset rng state to what it was.
+        _set_cuda_rng_state(orig_rng_state)
+
+    @contextlib.contextmanager
+    def fork(self, name=_MODEL_PARALLEL_RNG_TRACKER_NAME):
+        """Fork the cuda rng state, perform operations, and exit with
+        the original state."""
+        # Check if we have added the state
+        if name not in self.states_:
+            raise Exception("cuda rng state {} is not added".format(name))
+        # Store current rng state.
+        orig_cuda_rng_state = torch.cuda.get_rng_state()
+        # Set rng state to the desired one
+        _set_cuda_rng_state(self.states_[name])
+        # Do the stuff we wanted to do.
+        try:
+            yield
+        finally:
+            # Update the current rng state for later use.
+            self.states_[name] = torch.cuda.get_rng_state()
+            # And set the state to the original state we started with.
+            _set_cuda_rng_state(orig_cuda_rng_state)
+
+
+# RNG tracker object.
+_CUDA_RNG_STATE_TRACKER = CudaRNGStatesTracker()
+
+
+def get_cuda_rng_tracker():
+    """Get cuda rng tracker."""
+    return _CUDA_RNG_STATE_TRACKER
+
+
+def model_parallel_cuda_manual_seed(seed):
+    """Initialize model parallel cuda seed.
+
+    This function should be called after the model parallel is
+    initialized. Also, no torch.cuda.manual_seed should be called
+    after this function. Basically, this is replacement for that
+    function.
+    Two set of RNG states are tracked:
+        default state: This is for data parallelism and is the same among a
+                       set of model parallel GPUs but different across
+                       different model paralle groups. This is used for
+                       example for dropout in the non-tensor-model-parallel regions.
+        tensor-model-parallel state: This state is different among a set of model
+                              parallel GPUs, but the same across data parallel
+                              groups. This is used for example for dropout in
+                              model parallel regions.
+    """
+    # 2718 is just for fun and any POSITIVE value will work.
+    offset = seed + 2718
+    tensor_model_parallel_seed = offset + get_tensor_model_parallel_rank()
+    # Data parallel gets the original seed.
+    data_parallel_seed = seed
+
+    _CUDA_RNG_STATE_TRACKER.reset()
+    # Set the default state.
+    torch.cuda.manual_seed(data_parallel_seed)
+    # and model parallel state.
+    _CUDA_RNG_STATE_TRACKER.add(_MODEL_PARALLEL_RNG_TRACKER_NAME, tensor_model_parallel_seed)
+
+
+# TODO (mkozuki): Move the below gradient checkpoint related features to another (new) file.
+class CheckpointFunction(torch.autograd.Function):
+    """This function is adapted from torch.utils.checkpoint with
+       two main changes:
+           1) torch.cuda.set_rng_state is replaced with `_set_cuda_rng_state`
+           2) the states in the model parallel tracker are also properly
+              tracked/set/reset.
+    """
+
+    @staticmethod
+    def forward(ctx, run_function, *args):
+        ctx.run_function = run_function
+
+        # Copy the rng states.
+        ctx.fwd_cpu_rng_state = torch.get_rng_state()
+        ctx.fwd_cuda_rng_state = torch.cuda.get_rng_state()
+        ctx.fwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
+
+        with torch.no_grad():
+            outputs = run_function(*args)
+
+        # Divide hidden states across model parallel group and only keep
+        # the chunk corresponding to the current rank.
+        if _CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER is not None:
+            ctx.input_0_shape = args[0].data.shape
+            args[0].data = split_tensor_into_1d_equal_chunks(args[0].data)
+            args[0].data = _CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER.add(args[0].data)
+
+        # Store everything.
+        ctx.save_for_backward(*args)
+        return outputs
+
+    @staticmethod
+    def backward(ctx, *args):
+        if not torch.autograd._is_checkpoint_valid():
+            raise RuntimeError("Checkpointing is not compatible with .grad(), " "please use .backward() if possible")
+        inputs = ctx.saved_tensors
+        if _CHECKPOINTED_ACTIVATIONS_MEMORY_BUFFER is not None:
+            inputs[0].data = gather_split_1d_tensor(inputs[0].data)
+            inputs[0].data = inputs[0].data.view(ctx.input_0_shape)
+
+        # Store the current states.
+        bwd_cpu_rng_state = torch.get_rng_state()
+        bwd_cuda_rng_state = torch.cuda.get_rng_state()
+        bwd_cuda_rng_state_tracker = get_cuda_rng_tracker().get_states()
+
+        # Set the states to what it used to be before the forward pass.
+        torch.set_rng_state(ctx.fwd_cpu_rng_state)
+        _set_cuda_rng_state(ctx.fwd_cuda_rng_state)
+        get_cuda_rng_tracker().set_states(ctx.fwd_cuda_rng_state_tracker)
+
+        # Compute the forward pass.
+        detached_inputs = detach_variable(inputs)
+        with torch.enable_grad():
+            outputs = ctx.run_function(*detached_inputs)
+
+        # Set the states back to what it was at the start of this function.
+        torch.set_rng_state(bwd_cpu_rng_state)
+        _set_cuda_rng_state(bwd_cuda_rng_state)
+        get_cuda_rng_tracker().set_states(bwd_cuda_rng_state_tracker)
+
+        if isinstance(outputs, torch.Tensor):
+            outputs = (outputs,)
+        torch.autograd.backward(outputs, args)
+        grads = tuple(inp.grad if isinstance(inp, torch.Tensor) else inp for inp in detached_inputs)
+        return (None,) + grads
+
+
+def checkpoint(function, *args):
+    """Checkpoint a model or part of the model.
+    This has been directly copied from torch.utils.checkpoint."""
+    return CheckpointFunction.apply(function, *args)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/utils.py
new file mode 100644
index 0000000000..4c3362a494
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/tensor_parallel/utils.py
@@ -0,0 +1,54 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import torch
+
+from apex.transformer.utils import divide
+
+
+def split_tensor_along_last_dim(tensor, num_partitions, contiguous_split_chunks=False):
+    """Split a tensor along its last dimension.
+    Arguments:
+        tensor: input tensor.
+        num_partitions: number of partitions to split the tensor
+        contiguous_split_chunks: If True, make each chunk contiguous
+                                 in memory.
+    """
+    # Get the size and dimension.
+    last_dim = tensor.dim() - 1
+    last_dim_size = divide(tensor.size()[last_dim], num_partitions)
+    # Split.
+    tensor_list = torch.split(tensor, last_dim_size, dim=last_dim)
+    # Note: torch.split does not create contiguous tensors by default.
+    if contiguous_split_chunks:
+        return tuple(chunk.contiguous() for chunk in tensor_list)
+
+    return tensor_list
+
+
+class VocabUtility:
+    """Split the vocabulary into `world_size` chunks and return the
+    first and last index of the vocabulary belonging to the `rank`
+    partition: Note that indices in [fist, last)"""
+
+    @staticmethod
+    def vocab_range_from_per_partition_vocab_size(per_partition_vocab_size, rank, world_size):
+        index_f = rank * per_partition_vocab_size
+        index_l = index_f + per_partition_vocab_size
+        return index_f, index_l
+
+    @staticmethod
+    def vocab_range_from_global_vocab_size(global_vocab_size, rank, world_size):
+        per_partition_vocab_size = divide(global_vocab_size, world_size)
+        return VocabUtility.vocab_range_from_per_partition_vocab_size(per_partition_vocab_size, rank, world_size)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/arguments.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/arguments.py
new file mode 100644
index 0000000000..e142886678
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/arguments.py
@@ -0,0 +1,806 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+"""Megatron arguments."""
+
+import argparse
+import os
+
+import torch
+
+def parse_args(extra_args_provider=None, defaults={},
+               ignore_unknown_args=False):
+    """Parse all arguments."""
+    parser = argparse.ArgumentParser(description='Megatron-LM Arguments',
+                                     allow_abbrev=False)
+
+    # Standard arguments.
+    parser = _add_network_size_args(parser)
+    parser = _add_regularization_args(parser)
+    parser = _add_training_args(parser)
+    parser = _add_initialization_args(parser)
+    parser = _add_learning_rate_args(parser)
+    parser = _add_checkpointing_args(parser)
+    parser = _add_mixed_precision_args(parser)
+    parser = _add_distributed_args(parser)
+    parser = _add_validation_args(parser)
+    parser = _add_data_args(parser)
+    parser = _add_autoresume_args(parser)
+    parser = _add_biencoder_args(parser)
+    parser = _add_vit_args(parser)
+    parser = _add_logging_args(parser)
+
+    # Custom arguments.
+    if extra_args_provider is not None:
+        parser = extra_args_provider(parser)
+
+    # Parse.
+    if ignore_unknown_args:
+        args, _ = parser.parse_known_args()
+    else:
+        args = parser.parse_args()
+
+    # Distributed args.
+    args.rank = int(os.getenv('RANK', '0'))
+    args.world_size = int(os.getenv("WORLD_SIZE", '1'))
+    # Tensor model parallel size.
+    args.tensor_model_parallel_size = min(
+        args.tensor_model_parallel_size, args.world_size)
+    assert args.world_size % args.tensor_model_parallel_size == 0, 'world size'\
+        ' ({}) is not divisible by tensor model parallel size ({})'.format(
+            args.world_size, args.tensor_model_parallel_size)
+    # Pipeline model parallel size.
+    args.pipeline_model_parallel_size = min(
+        args.pipeline_model_parallel_size,
+        (args.world_size // args.tensor_model_parallel_size))
+    # Checks.
+    model_parallel_size = args.pipeline_model_parallel_size * \
+                          args.tensor_model_parallel_size
+    assert args.world_size % model_parallel_size == 0, 'world size is not'\
+        ' divisible by tensor parallel size ({}) times pipeline parallel ' \
+        'size ({})'.format(args.world_size, args.tensor_model_parallel_size,
+                           args.pipeline_model_parallel_size)
+    args.data_parallel_size = args.world_size // model_parallel_size
+    if args.rank == 0:
+        print('using world size: {}, data-parallel-size: {}, '
+              'tensor-model-parallel size: {}, '
+              'pipeline-model-parallel size: {} '.format(
+                  args.world_size, args.data_parallel_size,
+                  args.tensor_model_parallel_size,
+                  args.pipeline_model_parallel_size), flush=True)
+    if args.pipeline_model_parallel_size > 1:
+        if args.pipeline_model_parallel_split_rank is not None:
+            assert args.pipeline_model_parallel_split_rank < \
+                    args.pipeline_model_parallel_size, 'split rank needs'\
+                    ' to be less than pipeline model parallel size ({})'.format(
+                            args.pipeline_model_parallel_size)
+
+    # Deprecated arguments
+    assert args.batch_size is None, '--batch-size argument is no longer ' \
+        'valid, use --micro-batch-size instead'
+    del args.batch_size
+    assert args.warmup is None, '--warmup argument is no longer valid, use ' \
+        '--lr-warmup-fraction instead'
+    del args.warmup
+    assert args.model_parallel_size is None, '--model-parallel-size is no ' \
+        'longer valid, use --tensor-model-parallel-size instead'
+    del args.model_parallel_size
+    if args.checkpoint_activations:
+        args.activations_checkpoint_method = 'uniform'
+        if args.rank == 0:
+            print('--checkpoint-activations is no longer valid, '
+                  'use --activation-checkpoint-method instead. '
+                  'Defaulting to activation-checkpoint-method=uniform.')
+    del args.checkpoint_activations
+
+    # Set input defaults.
+    for key in defaults:
+        # For default to be valid, it should not be provided in the
+        # arguments that are passed to the program. We check this by
+        # ensuring the arg is set to None.
+        if getattr(args, key) is not None:
+            if args.rank == 0:
+                print('WARNING: overriding default arguments for {key}:{v} \
+                       with {key}:{v2}'.format(key=key, v=defaults[key],
+                                               v2=getattr(args, key)),
+                                               flush=True)
+        else:
+            setattr(args, key, defaults[key])
+
+    # Batch size.
+    assert args.micro_batch_size is not None
+    assert args.micro_batch_size > 0
+    if args.global_batch_size is None:
+        args.global_batch_size = args.micro_batch_size * args.data_parallel_size
+        if args.rank == 0:
+            print('setting global batch size to {}'.format(
+                args.global_batch_size), flush=True)
+    assert args.global_batch_size > 0
+    if args.num_layers_per_virtual_pipeline_stage is not None:
+        assert args.pipeline_model_parallel_size > 2, \
+            'pipeline-model-parallel size should be greater than 2 with ' \
+            'interleaved schedule'
+        assert args.num_layers % args.num_layers_per_virtual_pipeline_stage == 0, \
+            'number of layers is not divisible by number of layers per virtual ' \
+            'pipeline stage'
+        args.virtual_pipeline_model_parallel_size = \
+            (args.num_layers // args.pipeline_model_parallel_size) // \
+            args.num_layers_per_virtual_pipeline_stage
+    else:
+        args.virtual_pipeline_model_parallel_size = None
+
+    # Parameters dtype.
+    args.params_dtype = torch.float
+    if args.fp16:
+        assert not args.bf16
+        args.params_dtype = torch.half
+    if args.bf16:
+        assert not args.fp16
+        args.params_dtype = torch.bfloat16
+        # bfloat16 requires gradient accumulation and all-reduce to
+        # be done in fp32.
+        if not args.accumulate_allreduce_grads_in_fp32:
+            args.accumulate_allreduce_grads_in_fp32 = True
+            if args.rank == 0:
+                print('accumulate and all-reduce gradients in fp32 for '
+                      'bfloat16 data type.', flush=True)
+
+    if args.rank == 0:
+        print('using {} for parameters ...'.format(args.params_dtype),
+              flush=True)
+
+    # If we do accumulation and all-reduces in fp32, we need to have local DDP
+    # and we should make sure use-contiguous-buffers-in-local-ddp is not off.
+    if args.accumulate_allreduce_grads_in_fp32:
+        assert args.DDP_impl == 'local'
+        assert args.use_contiguous_buffers_in_local_ddp
+
+    # For torch DDP, we do not use contiguous buffer
+    if args.DDP_impl == 'torch':
+        args.use_contiguous_buffers_in_local_ddp = False
+
+    if args.dataloader_type is None:
+        args.dataloader_type = 'single'
+
+    # Consumed tokens.
+    args.consumed_train_samples = 0
+    args.consumed_valid_samples = 0
+
+    # Iteration-based training.
+    if args.train_iters:
+        # If we use iteration-based training, make sure the
+        # sample-based options are off.
+        assert args.train_samples is None, \
+            'expected iteration-based training'
+        assert args.lr_decay_samples is None, \
+            'expected iteration-based learning rate decay'
+        assert args.lr_warmup_samples == 0, \
+            'expected iteration-based learning rate warmup'
+        assert args.rampup_batch_size is None, \
+            'expected no batch-size rampup for iteration-based training'
+        if args.lr_warmup_fraction is not None:
+            assert args.lr_warmup_iters == 0, \
+                'can only specify one of lr-warmup-fraction and lr-warmup-iters'
+
+    # Sample-based training.
+    if args.train_samples:
+        # If we use sample-based training, make sure the
+        # iteration-based options are off.
+        assert args.train_iters is None, \
+            'expected sample-based training'
+        assert args.lr_decay_iters is None, \
+            'expected sample-based learning rate decay'
+        assert args.lr_warmup_iters == 0, \
+            'expected sample-based learnig rate warmup'
+        if args.lr_warmup_fraction is not None:
+            assert args.lr_warmup_samples == 0, \
+                'can only specify one of lr-warmup-fraction ' \
+                'and lr-warmup-samples'
+
+    # Check required arguments.
+    required_args = ['num_layers', 'hidden_size', 'num_attention_heads',
+                     'max_position_embeddings']
+    for req_arg in required_args:
+        _check_arg_is_not_none(args, req_arg)
+
+    # Checks.
+    if args.ffn_hidden_size is None:
+        args.ffn_hidden_size = 4 * args.hidden_size
+
+    if args.kv_channels is None:
+        assert args.hidden_size % args.num_attention_heads == 0
+        args.kv_channels = args.hidden_size // args.num_attention_heads
+
+    if args.seq_length is not None:
+        assert args.encoder_seq_length is None
+        args.encoder_seq_length = args.seq_length
+    else:
+        assert args.encoder_seq_length is not None
+        args.seq_length = args.encoder_seq_length
+
+    if args.seq_length is not None:
+        assert args.max_position_embeddings >= args.seq_length
+    if args.decoder_seq_length is not None:
+        assert args.max_position_embeddings >= args.decoder_seq_length
+    if args.lr is not None:
+        assert args.min_lr <= args.lr
+    if args.save is not None:
+        assert args.save_interval is not None
+    # Mixed precision checks.
+    if args.fp16_lm_cross_entropy:
+        assert args.fp16, 'lm cross entropy in fp16 only support in fp16 mode.'
+    if args.fp32_residual_connection:
+        assert args.fp16 or args.bf16, \
+            'residual connection in fp32 only supported when using fp16 or bf16.'
+    # Activation checkpointing.
+    if args.distribute_checkpointed_activations:
+        assert args.tensor_model_parallel_size > 1, 'can distribute ' \
+            'checkpointed activations only across tensor model ' \
+            'parallel groups'
+        assert args.activations_checkpoint_method is not None, \
+            'for distribute-checkpointed-activations to work you '\
+            'need to use a activation-checkpoint method '
+        assert args.num_layers_per_virtual_pipeline_stage is None, \
+            'currently distrobuted checkpoint activations only supported for ' \
+            'nointerleaved pipeline parallelism'
+
+    _print_args(args)
+    return args
+
+
+def _print_args(args):
+    """Print arguments."""
+    if args.rank == 0:
+        print('------------------------ arguments ------------------------',
+              flush=True)
+        str_list = []
+        for arg in vars(args):
+            dots = '.' * (48 - len(arg))
+            str_list.append('  {} {} {}'.format(arg, dots, getattr(args, arg)))
+        for arg in sorted(str_list, key=lambda x: x.lower()):
+            print(arg, flush=True)
+        print('-------------------- end of arguments ---------------------',
+              flush=True)
+
+
+def _check_arg_is_not_none(args, arg):
+    assert getattr(args, arg) is not None, '{} argument is None'.format(arg)
+
+
+def _add_network_size_args(parser):
+    group = parser.add_argument_group(title='network size')
+
+    group.add_argument('--num-layers', type=int, default=None,
+                       help='Number of transformer layers.')
+    group.add_argument('--hidden-size', type=int, default=None,
+                       help='Tansformer hidden size.')
+    group.add_argument('--ffn-hidden-size', type=int, default=None,
+                       help='Transformer Feed-Forward Network hidden size. '
+                       'This is set to 4*hidden-size if not provided')
+    group.add_argument('--num-attention-heads', type=int, default=None,
+                       help='Number of transformer attention heads.')
+    group.add_argument('--kv-channels', type=int, default=None,
+                       help='Projection weights dimension in multi-head '
+                       'attention. This is set to '
+                       '   args.hidden_size // args.num_attention_heads '
+                       'if not provided.')
+    group.add_argument('--max-position-embeddings', type=int, default=None,
+                       help='Maximum number of position embeddings to use. '
+                       'This is the size of position embedding.')
+    group.add_argument('--make-vocab-size-divisible-by', type=int, default=128,
+                       help='Pad the vocab size to be divisible by this value.'
+                       'This is added for computational efficieny reasons.')
+    group.add_argument('--layernorm-epsilon', type=float, default=1e-5,
+                       help='Layer norm epsilon.')
+    group.add_argument('--apply-residual-connection-post-layernorm',
+                       action='store_true',
+                       help='If set, use original BERT residula connection '
+                       'ordering.')
+    group.add_argument('--openai-gelu', action='store_true',
+                       help='Use OpenAIs GeLU implementation. This option'
+                       'should not be used unless for backward compatibility'
+                       'reasons.')
+    group.add_argument('--onnx-safe', type=bool, required=False,
+                       help='Use workarounds for known problems with '
+                       'Torch ONNX exporter')
+    group.add_argument('--bert-no-binary-head', action='store_false',
+                       help='Disable BERT binary head.',
+                       dest='bert_binary_head')
+
+    return parser
+
+
+def _add_logging_args(parser):
+    group = parser.add_argument_group(title='logging')
+
+    group.add_argument('--log-params-norm', action='store_true',
+                       help='If set, calculate and log parameters norm.')
+    group.add_argument('--log-num-zeros-in-grad', action='store_true',
+                       help='If set, calculate and log the number of zeros in gradient.')
+    group.add_argument('--tensorboard-log-interval', type=int, default=1,
+                       help='Report to tensorboard interval.')
+    group.add_argument('--tensorboard-queue-size', type=int, default=1000,
+                       help='Size of the tensorboard queue for pending events '
+                       'and summaries before one of the ‘add’ calls forces a '
+                       'flush to disk.')
+    group.add_argument('--log-timers-to-tensorboard', action='store_true',
+                       help='If set, write timers to tensorboard.')
+    group.add_argument('--log-batch-size-to-tensorboard', action='store_true',
+                       help='If set, write batch-size to tensorboard.')
+    group.add_argument('--no-log-learnig-rate-to-tensorboard',
+                       action='store_false',
+                       help='Disable learning rate logging to tensorboard.',
+                       dest='log_learning_rate_to_tensorboard')
+    group.add_argument('--no-log-loss-scale-to-tensorboard',
+                       action='store_false',
+                       help='Disable loss-scale logging to tensorboard.',
+                       dest='log_loss_scale_to_tensorboard')
+    group.add_argument('--log-validation-ppl-to-tensorboard',
+                       action='store_true',
+                       help='If set, write validation perplexity to '
+                       'tensorboard.')
+    group.add_argument('--log-memory-to-tensorboard',
+                       action='store_true',
+                       help='Enable memory logging to tensorboard.')
+
+    return parser
+
+
+def _add_regularization_args(parser):
+    group = parser.add_argument_group(title='regularization')
+
+    group.add_argument('--attention-dropout', type=float, default=0.1,
+                       help='Post attention dropout probability.')
+    group.add_argument('--hidden-dropout', type=float, default=0.1,
+                       help='Dropout probability for hidden state transformer.')
+    group.add_argument('--weight-decay', type=float, default=0.01,
+                       help='Weight decay coefficient for L2 regularization.')
+    group.add_argument('--clip-grad', type=float, default=1.0,
+                       help='Gradient clipping based on global L2 norm.')
+    group.add_argument('--adam-beta1', type=float, default=0.9,
+                       help='First coefficient for computing running averages '
+                       'of gradient and its square')
+    group.add_argument('--adam-beta2', type=float, default=0.999,
+                       help='Second coefficient for computing running averages '
+                       'of gradient and its square')
+    group.add_argument('--adam-eps', type=float, default=1e-08,
+                       help='Term added to the denominator to improve'
+                       'numerical stability')
+    group.add_argument('--sgd-momentum', type=float, default=0.9,
+                       help='Momentum factor for sgd')
+
+    return parser
+
+
+def _add_training_args(parser):
+    group = parser.add_argument_group(title='training')
+
+    group.add_argument('--micro-batch-size', type=int, default=None,
+                       help='Batch size per model instance (local batch size). '
+                       'Global batch size is local batch size times data '
+                       'parallel size times number of micro batches.')
+    group.add_argument('--batch-size', type=int, default=None,
+                       help='Old batch size parameter, do not use. '
+                       'Use --micro-batch-size instead')
+    group.add_argument('--global-batch-size', type=int, default=None,
+                       help='Training batch size. If set, it should be a '
+                       'multiple of micro-batch-size times data-parallel-size. '
+                       'If this value is None, then '
+                       'use micro-batch-size * data-parallel-size as the '
+                       'global batch size. This choice will result in 1 for '
+                       'number of micro-batches.')
+    group.add_argument('--rampup-batch-size', nargs='*', default=None,
+                       help='Batch size ramp up with the following values:'
+                       '  --rampup-batch-size <start batch size> '
+                       '                      <batch size incerement> '
+                       '                      <ramp-up samples> '
+                       'For example:'
+                       '   --rampup-batch-size 16 8 300000 \ '
+                       '   --global-batch-size 1024'
+                       'will start with global batch size 16 and over '
+                       ' (1024 - 16) / 8 = 126 intervals will increase'
+                       'the batch size linearly to 1024. In each interval'
+                       'we will use approximately 300000 / 126 = 2380 samples.')
+    group.add_argument('--checkpoint-activations', action='store_true',
+                       help='Checkpoint activation to allow for training '
+                       'with larger models, sequences, and batch sizes.')
+    group.add_argument('--distribute-checkpointed-activations',
+                       action='store_true',
+                       help='If set, distribute checkpointed activations '
+                       'across model parallel group.')
+    group.add_argument('--activations-checkpoint-method', type=str, default=None,
+                       choices=['uniform', 'block'],
+                       help='1) uniform: uniformly divide the total number of '
+                       'Transformer layers and checkpoint the input activation of '
+                       'each divided chunk, '
+                       '2) checkpoint the input activations of only a set number of '
+                       'individual Transformer layers per pipeline stage and do the '
+                       'rest without any checkpointing'
+                       'default) do not apply activations checkpoint to any layers')
+    group.add_argument('--activations-checkpoint-num-layers', type=int, default=1,
+                       help='1) uniform: the number of Transformer layers in each '
+                       'uniformly divided checkpoint unit, '
+                       '2) block: the number of individual Transformer layers '
+                       'to checkpoint within each pipeline stage.')
+    group.add_argument('--train-iters', type=int, default=None,
+                       help='Total number of iterations to train over all '
+                       'training runs. Note that either train-iters or '
+                       'train-samples should be provided.')
+    group.add_argument('--train-samples', type=int, default=None,
+                       help='Total number of samples to train over all '
+                       'training runs. Note that either train-iters or '
+                       'train-samples should be provided.')
+    group.add_argument('--log-interval', type=int, default=100,
+                       help='Report loss and timing interval.')
+    group.add_argument('--exit-interval', type=int, default=None,
+                       help='Exit the program after the iteration is divisible '
+                       'by this value.')
+    group.add_argument('--exit-duration-in-mins', type=int, default=None,
+                       help='Exit the program after this many minutes.')
+    group.add_argument('--tensorboard-dir', type=str, default=None,
+                       help='Write TensorBoard logs to this directory.')
+    group.add_argument('--no-masked-softmax-fusion',
+                       action='store_false',
+                       help='Disable fusion of query_key_value scaling, '
+                       'masking, and softmax.',
+                       dest='masked_softmax_fusion')
+    group.add_argument('--no-bias-gelu-fusion', action='store_false',
+                       help='Disable bias and gelu fusion.',
+                       dest='bias_gelu_fusion')
+    group.add_argument('--no-bias-dropout-fusion', action='store_false',
+                       help='Disable bias and dropout fusion.',
+                       dest='bias_dropout_fusion')
+    group.add_argument('--optimizer', type=str, default='adam',
+                       choices=['adam', 'sgd'],
+                       help='Optimizer function')
+    group.add_argument('--dataloader-type', type=str, default=None,
+                       choices=['single', 'cyclic'],
+                       help='Single pass vs multiple pass data loader')
+    group.add_argument('--no-async-tensor-model-parallel-allreduce',
+                       action='store_true',
+                       help='Disable asynchronous execution of '
+                       'tensor-model-parallel all-reduce with weight '
+                       'gradient compuation of a column-linear layer.')
+    return parser
+
+
+def _add_initialization_args(parser):
+    group = parser.add_argument_group(title='initialization')
+
+    group.add_argument('--seed', type=int, default=1234,
+                       help='Random seed used for python, numpy, '
+                       'pytorch, and cuda.')
+    group.add_argument('--init-method-std', type=float, default=0.02,
+                       help='Standard deviation of the zero mean normal '
+                       'distribution used for weight initialization.')
+    group.add_argument('--init-method-xavier-uniform', action='store_true',
+                       help='Enable Xavier uniform parameter initialization')
+
+    return parser
+
+
+def _add_learning_rate_args(parser):
+    group = parser.add_argument_group(title='learning rate')
+
+    group.add_argument('--lr', type=float, default=None,
+                       help='Initial learning rate. Depending on decay style '
+                       'and initial warmup, the learing rate at each '
+                       'iteration would be different.')
+    group.add_argument('--lr-decay-style', type=str, default='linear',
+                       choices=['constant', 'linear', 'cosine'],
+                       help='Learning rate decay function.')
+    group.add_argument('--lr-decay-iters', type=int, default=None,
+                       help='number of iterations to decay learning rate over,'
+                       ' If None defaults to `--train-iters`')
+    group.add_argument('--lr-decay-samples', type=int, default=None,
+                       help='number of samples to decay learning rate over,'
+                       ' If None defaults to `--train-samples`')
+    group.add_argument('--lr-warmup-fraction', type=float, default=None,
+                       help='fraction of lr-warmup-(iters/samples) to use '
+                       'for warmup (as a float)')
+    group.add_argument('--lr-warmup-iters', type=int, default=0,
+                       help='number of iterations to linearly warmup '
+                       'learning rate over.')
+    group.add_argument('--lr-warmup-samples', type=int, default=0,
+                       help='number of samples to linearly warmup '
+                       'learning rate over.')
+    group.add_argument('--warmup', type=int, default=None,
+                       help='Old lr warmup argument, do not use. Use one of the'
+                       '--lr-warmup-* arguments above')
+    group.add_argument('--min-lr', type=float, default=0.0,
+                       help='Minumum value for learning rate. The scheduler'
+                       'clip values below this threshold.')
+    group.add_argument('--override-lr-scheduler', action='store_true',
+                       help='Reset the values of the scheduler (learning rate,'
+                       'warmup iterations, minimum learning rate, maximum '
+                       'number of iterations, and decay style from input '
+                       'arguments and ignore values from checkpoints. Note'
+                       'that all the above values will be reset.')
+    group.add_argument('--use-checkpoint-lr-scheduler', action='store_true',
+                       help='Use checkpoint to set the values of the scheduler '
+                       '(learning rate, warmup iterations, minimum learning '
+                       'rate, maximum number of iterations, and decay style '
+                       'from checkpoint and ignore input arguments.')
+
+    return parser
+
+
+def _add_checkpointing_args(parser):
+    group = parser.add_argument_group(title='checkpointing')
+
+    group.add_argument('--save', type=str, default=None,
+                       help='Output directory to save checkpoints to.')
+    group.add_argument('--save-interval', type=int, default=None,
+                       help='Number of iterations between checkpoint saves.')
+    group.add_argument('--no-save-optim', action='store_true', default=None,
+                       help='Do not save current optimizer.')
+    group.add_argument('--no-save-rng', action='store_true', default=None,
+                       help='Do not save current rng state.')
+    group.add_argument('--load', type=str, default=None,
+                       help='Directory containing a model checkpoint.')
+    group.add_argument('--no-load-optim', action='store_true', default=None,
+                       help='Do not load optimizer when loading checkpoint.')
+    group.add_argument('--no-load-rng', action='store_true', default=None,
+                       help='Do not load rng state when loading checkpoint.')
+    group.add_argument('--finetune', action='store_true',
+                       help='Load model for finetuning. Do not load optimizer '
+                       'or rng state from checkpoint and set iteration to 0. '
+                       'Assumed when loading a release checkpoint.')
+
+    return parser
+
+
+def _add_mixed_precision_args(parser):
+    group = parser.add_argument_group(title='mixed precision')
+
+    group.add_argument('--fp16', action='store_true',
+                       help='Run model in fp16 mode.')
+    group.add_argument('--bf16', action='store_true',
+                       help='Run model in bfloat16 mode.')
+    group.add_argument('--loss-scale', type=float, default=None,
+                       help='Static loss scaling, positive power of 2 '
+                       'values can improve fp16 convergence. If None, dynamic'
+                       'loss scaling is used.')
+    group.add_argument('--initial-loss-scale', type=float, default=2**32,
+                       help='Initial loss-scale for dynamic loss scaling.')
+    group.add_argument('--min-loss-scale', type=float, default=1.0,
+                       help='Minimum loss scale for dynamic loss scale.')
+    group.add_argument('--loss-scale-window', type=float, default=1000,
+                       help='Window over which to raise/lower dynamic scale.')
+    group.add_argument('--hysteresis', type=int, default=2,
+                       help='hysteresis for dynamic loss scaling')
+    group.add_argument('--fp32-residual-connection', action='store_true',
+                       help='Move residual connections to fp32.')
+    group.add_argument('--no-query-key-layer-scaling', action='store_false',
+                       help='Do not scale Q * K^T by 1 / layer-number.',
+                       dest='apply_query_key_layer_scaling')
+    group.add_argument('--attention-softmax-in-fp32', action='store_true',
+                       help='Run attention masking and softmax in fp32. '
+                       'This flag is ignored unless '
+                       '--no-query-key-layer-scaling is specified.')
+    group.add_argument('--accumulate-allreduce-grads-in-fp32',
+                       action='store_true',
+                       help='Gradient accumulation and all-reduce in fp32.')
+    group.add_argument('--fp16-lm-cross-entropy', action='store_true',
+                       help='Move the cross entropy unreduced loss calculation'
+                       'for lm head to fp16.')
+
+    return parser
+
+
+def _add_distributed_args(parser):
+    group = parser.add_argument_group(title='distributed')
+
+    group.add_argument('--tensor-model-parallel-size', type=int, default=1,
+                       help='Degree of tensor model parallelism.')
+    group.add_argument('--pipeline-model-parallel-size', type=int, default=1,
+                       help='Degree of pipeline model parallelism.')
+    group.add_argument('--pipeline-model-parallel-split-rank',
+                       type=int, default=None,
+                       help='Rank where encoder and decoder should be split.')
+    group.add_argument('--model-parallel-size', type=int, default=None,
+                       help='Old model parallel argument, do not use. Use '
+                       '--tensor-model-parallel-size instead.')
+    group.add_argument('--num-layers-per-virtual-pipeline-stage', type=int, default=None,
+                       help='Number of layers per virtual pipeline stage')
+    group.add_argument('--distributed-backend', default='nccl',
+                       choices=['nccl', 'gloo'],
+                       help='Which backend to use for distributed training.')
+    group.add_argument('--DDP-impl', default='local',
+                       choices=['local', 'torch'],
+                       help='which DistributedDataParallel implementation '
+                       'to use.')
+    group.add_argument('--no-contiguous-buffers-in-local-ddp',
+                       action='store_false', help='If set, dont use '
+                       'contiguous buffer in local DDP.',
+                       dest='use_contiguous_buffers_in_local_ddp')
+    group.add_argument('--no-scatter-gather-tensors-in-pipeline', action='store_false',
+                       help='Use scatter/gather to optimize communication of tensors in pipeline',
+                       dest='scatter_gather_tensors_in_pipeline')
+    group.add_argument('--local_rank', type=int, default=None,
+                       help='local rank passed from distributed launcher.')
+    group.add_argument('--lazy-mpu-init', type=bool, required=False,
+                       help='If set to True, initialize_megatron() '
+                       'skips DDP initialization and returns function to '
+                       'complete it instead.Also turns on '
+                       '--use-cpu-initialization flag. This is for '
+                       'external DDP manager.' )
+    group.add_argument('--use-cpu-initialization', action='store_true',
+                       default=None, help='If set, affine parallel weights '
+                       'initialization uses CPU' )
+    group.add_argument('--empty-unused-memory-level', default=0, type=int,
+                       choices=[0, 1, 2],
+                       help='Call torch.cuda.empty_cache() each iteration '
+                       '(training and eval), to reduce fragmentation.'
+                       '0=off, 1=moderate, 2=aggressive.')
+    return parser
+
+
+def _add_validation_args(parser):
+    group = parser.add_argument_group(title='validation')
+
+    group.add_argument('--eval-iters', type=int, default=100,
+                       help='Number of iterations to run for evaluation'
+                       'validation/test for.')
+    group.add_argument('--eval-interval', type=int, default=1000,
+                       help='Interval between running evaluation on '
+                       'validation set.')
+
+    return parser
+
+
+def _add_data_args(parser):
+    group = parser.add_argument_group(title='data and dataloader')
+
+    group.add_argument('--data-path', nargs='*', default=None,
+                       help='Path to the training dataset. Accepted format:'
+                       '1) a single data path, 2) multiple datasets in the'
+                       'form: dataset1-weight dataset1-path dataset2-weight '
+                       'dataset2-path ...')
+    group.add_argument('--split', type=str, default='969, 30, 1',
+                       help='Comma-separated list of proportions for training,'
+                       ' validation, and test split. For example the split '
+                       '`90,5,5` will use 90%% of data for training, 5%% for '
+                       'validation and 5%% for test.')
+    group.add_argument('--vocab-file', type=str, default=None,
+                       help='Path to the vocab file.')
+    group.add_argument('--merge-file', type=str, default=None,
+                       help='Path to the BPE merge file.')
+    group.add_argument('--vocab-extra-ids', type=int, default=0,
+                       help='Number of additional vocabulary tokens. '
+                            'They are used for span masking in the T5 model')
+    group.add_argument('--seq-length', type=int, default=None,
+                       help='Maximum sequence length to process.')
+    group.add_argument('--encoder-seq-length', type=int, default=None,
+                       help='Maximum encoder sequence length to process.'
+                       'This should be exclusive of --seq-length')
+    group.add_argument('--decoder-seq-length', type=int, default=None,
+                       help="Maximum decoder sequence length to process.")
+    group.add_argument('--retriever-seq-length', type=int, default=256,
+                       help='Maximum sequence length for the biencoder model '
+                        ' for retriever')
+    group.add_argument('--sample-rate', type=float, default=1.0,
+                       help='sample rate for training data. Supposed to be 0 '
+                            ' < sample_rate < 1')
+    group.add_argument('--mask-prob', type=float, default=0.15,
+                       help='Probability of replacing a token with mask.')
+    group.add_argument('--short-seq-prob', type=float, default=0.1,
+                       help='Probability of producing a short sequence.')
+    group.add_argument('--mmap-warmup', action='store_true',
+                       help='Warm up mmap files.')
+    group.add_argument('--num-workers', type=int, default=2,
+                       help="Dataloader number of workers.")
+    group.add_argument('--tokenizer-type', type=str,
+                       default=None,
+                       choices=['BertWordPieceLowerCase',
+                                'BertWordPieceCase',
+                                'GPT2BPETokenizer'],
+                       help='What type of tokenizer to use.')
+    group.add_argument('--data-impl', type=str, default='infer',
+                       choices=['lazy', 'cached', 'mmap', 'infer'],
+                       help='Implementation of indexed datasets.')
+    group.add_argument('--reset-position-ids', action='store_true',
+                       help='Reset posistion ids after end-of-document token.')
+    group.add_argument('--reset-attention-mask', action='store_true',
+                       help='Reset self attention maske after '
+                       'end-of-document token.')
+    group.add_argument('--eod-mask-loss', action='store_true',
+                       help='Mask loss for the end of document tokens.')
+
+    return parser
+
+
+def _add_autoresume_args(parser):
+    group = parser.add_argument_group(title='autoresume')
+
+    group.add_argument('--adlr-autoresume', action='store_true',
+                       help='Enable autoresume on adlr cluster.')
+    group.add_argument('--adlr-autoresume-interval', type=int, default=1000,
+                       help='Intervals over which check for autoresume'
+                       'termination signal')
+
+    return parser
+
+
+def _add_biencoder_args(parser):
+    group = parser.add_argument_group(title='biencoder')
+
+    # network size
+    group.add_argument('--ict-head-size', type=int, default=None,
+                       help='Size of block embeddings to be used in ICT and '
+                        'REALM (paper default: 128)')
+    group.add_argument('--biencoder-projection-dim', type=int, default=0,
+                       help='Size of projection head used in biencoder (paper'
+                        ' default: 128)')
+    group.add_argument('--biencoder-shared-query-context-model', action='store_true',
+                        help='Whether to share the parameters of the query '
+                        'and context models or not')
+
+    # checkpointing
+    group.add_argument('--ict-load', type=str, default=None,
+                       help='Directory containing an ICTBertModel checkpoint')
+    group.add_argument('--bert-load', type=str, default=None,
+                       help='Directory containing an BertModel checkpoint '
+                       '(needed to start ICT and REALM)')
+
+    # data
+    group.add_argument('--titles-data-path', type=str, default=None,
+                       help='Path to titles dataset used for ICT')
+    group.add_argument('--query-in-block-prob', type=float, default=0.1,
+                       help='Probability of keeping query in block for '
+                       'ICT dataset')
+    group.add_argument('--use-one-sent-docs', action='store_true',
+                       help='Whether to use one sentence documents in ICT')
+    group.add_argument('--evidence-data-path', type=str, default=None,
+                       help='Path to Wikipedia Evidence frm DPR paper')
+
+    # training
+    group.add_argument('--retriever-report-topk-accuracies', nargs='+', type=int,
+                        default=[], help="Which top-k accuracies to report "
+                        "(e.g. '1 5 20')")
+    group.add_argument('--retriever-score-scaling', action='store_true',
+                       help='Whether to scale retriever scores by inverse '
+                        'square root of hidden size')
+
+    # faiss index
+    group.add_argument('--block-data-path', type=str, default=None,
+                       help='Where to save/load BlockData to/from')
+    group.add_argument('--embedding-path', type=str, default=None,
+                       help='Where to save/load Open-Retrieval Embedding'
+                        ' data to/from')
+
+    # indexer
+    group.add_argument('--indexer-batch-size', type=int, default=128,
+                       help='How large of batches to use when doing indexing '
+                       'jobs')
+    group.add_argument('--indexer-log-interval', type=int, default=1000,
+                       help='After how many batches should the indexer '
+                       'report progress')
+    return parser
+
+
+def _add_vit_args(parser):
+    group = parser.add_argument_group(title="vit")
+
+    group.add_argument('--num-classes', type=int, default=1000,
+                       help='num of classes in vision classificaiton task')
+    group.add_argument('--img-dim', type=int, default=224,
+                       help='Image size for vision classification task')
+    group.add_argument('--num-channels', type=int, default=3,
+                       help='Number of channels in input image data')
+    group.add_argument('--patch-dim', type=int, default=16,
+                       help='patch dimension used in vit')
+
+    return parser
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/commons.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/commons.py
new file mode 100644
index 0000000000..4c97b4420e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/commons.py
@@ -0,0 +1,87 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import os
+import random
+
+import numpy
+import torch
+
+from apex import transformer
+from apex.transformer.testing import global_vars
+
+
+TEST_SUCCESS_MESSAGE = ">> passed the test :-)"
+
+
+class IdentityLayer(torch.nn.Module):
+    def __init__(self, size, scale=1.0):
+        super(IdentityLayer, self).__init__()
+        self.weight = torch.nn.Parameter(scale * torch.randn(size))
+
+    def forward(self):
+        return self.weight
+
+
+def set_random_seed(seed):
+    """Set random seed for reproducibility."""
+    random.seed(seed)
+    numpy.random.seed(seed)
+    torch.manual_seed(seed)
+    transformer.tensor_parallel.model_parallel_cuda_manual_seed(seed)
+
+
+def initialize_distributed(backend='nccl'):
+    """Initialize torch.distributed."""
+    # Get local rank in case it is provided.
+    # parser = argparse.ArgumentParser()
+    # parser.add_argument('--local_rank', type=int, default=None,
+    #                    help='local rank passed from distributed launcher')
+    # args = parser.parse_args()
+    args = global_vars.get_args()
+    local_rank = args.local_rank
+
+    # Get rank and world size.
+    rank = int(os.getenv('RANK', '0'))
+    world_size = int(os.getenv("WORLD_SIZE", '1'))
+
+    print('> initializing torch.distributed with local rank: {}, '
+          'rank: {}, world size: {}'.format(local_rank, rank, world_size))
+
+    # Set the device id.
+    device = rank % torch.cuda.device_count()
+    if local_rank is not None:
+        device = local_rank
+    torch.cuda.set_device(device)
+
+    # Call the init process.
+    init_method = 'tcp://'
+    master_ip = os.getenv('MASTER_ADDR', 'localhost')
+    master_port = os.getenv('MASTER_PORT', '6000')
+    init_method += master_ip + ':' + master_port
+    torch.distributed.init_process_group(
+        backend=backend,
+        world_size=world_size,
+        rank=rank,
+        init_method=init_method)
+
+
+def print_separator(message):
+    torch.distributed.barrier()
+    filler_len = (78 - len(message)) // 2
+    filler = '-' * filler_len
+    string = '\n' + filler + ' {} '.format(message) + filler
+    if torch.distributed.get_rank() == 0:
+        print(string, flush=True)
+    torch.distributed.barrier()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/global_vars.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/global_vars.py
new file mode 100644
index 0000000000..6a46ada078
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/global_vars.py
@@ -0,0 +1,260 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+"""Megatron global variables."""
+import os
+import sys
+import time
+
+import torch
+
+from apex.transformer.microbatches import build_num_microbatches_calculator
+from .arguments import parse_args
+
+_GLOBAL_ARGS = None
+_GLOBAL_NUM_MICROBATCHES_CALCULATOR = None
+_GLOBAL_TOKENIZER = None
+_GLOBAL_TENSORBOARD_WRITER = None
+_GLOBAL_ADLR_AUTORESUME = None
+_GLOBAL_TIMERS = None
+
+
+def get_args():
+    """Return arguments."""
+    _ensure_var_is_initialized(_GLOBAL_ARGS, 'args')
+    return _GLOBAL_ARGS
+
+
+def get_num_microbatches():
+    return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get()
+
+
+def get_current_global_batch_size():
+    return _GLOBAL_NUM_MICROBATCHES_CALCULATOR.get_current_global_batch_size()
+
+
+def update_num_microbatches(consumed_samples, consistency_check=True):
+    _GLOBAL_NUM_MICROBATCHES_CALCULATOR.update(consumed_samples,
+                                               consistency_check)
+
+
+# def get_tokenizer():
+#     """Return tokenizer."""
+#     _ensure_var_is_initialized(_GLOBAL_TOKENIZER, 'tokenizer')
+#     return _GLOBAL_TOKENIZER
+
+
+def get_tensorboard_writer():
+    """Return tensorboard writer. It can be None so no need
+    to check if it is initialized."""
+    return _GLOBAL_TENSORBOARD_WRITER
+
+
+def get_adlr_autoresume():
+    """ADLR autoresume object. It can be None so no need
+    to check if it is initialized."""
+    return _GLOBAL_ADLR_AUTORESUME
+
+
+def get_timers():
+    """Return timers."""
+    _ensure_var_is_initialized(_GLOBAL_TIMERS, 'timers')
+    return _GLOBAL_TIMERS
+
+
+def set_global_variables(extra_args_provider=None, args_defaults={},
+                         ignore_unknown_args=False):
+    """Set args, tokenizer, tensorboard-writer, adlr-autoresume, and timers."""
+    args = _parse_args(extra_args_provider=extra_args_provider,
+                       defaults=args_defaults,
+                       ignore_unknown_args=ignore_unknown_args)
+    # _build_num_microbatches_calculator(args)
+    # if args.vocab_file:
+    #     _ = _build_tokenizer(args)
+    _set_tensorboard_writer(args)
+    _set_adlr_autoresume(args)
+    _set_timers()
+
+
+def _parse_args(extra_args_provider=None, defaults={},
+                ignore_unknown_args=False):
+    """Parse entire arguments."""
+    global _GLOBAL_ARGS
+    _ensure_var_is_not_initialized(_GLOBAL_ARGS, 'args')
+    _GLOBAL_ARGS = parse_args(extra_args_provider=extra_args_provider,
+                              defaults=defaults,
+                              ignore_unknown_args=ignore_unknown_args)
+    return _GLOBAL_ARGS
+
+
+def _build_num_microbatches_calculator(args):
+
+    global _GLOBAL_NUM_MICROBATCHES_CALCULATOR
+    _ensure_var_is_not_initialized(_GLOBAL_NUM_MICROBATCHES_CALCULATOR,
+                                   'num microbatches calculator')
+
+    _GLOBAL_NUM_MICROBATCHES_CALCULATOR = build_num_microbatches_calculator(
+        args)
+
+
+# def _build_tokenizer(args):
+#     """Initialize tokenizer."""
+#     global _GLOBAL_TOKENIZER
+#     _ensure_var_is_not_initialized(_GLOBAL_TOKENIZER, 'tokenizer')
+#     _GLOBAL_TOKENIZER = build_tokenizer(args)
+#     return _GLOBAL_TOKENIZER
+
+
+# def rebuild_tokenizer(args):
+#     global _GLOBAL_TOKENIZER
+#     _GLOBAL_TOKENIZER = None
+#     return _build_tokenizer(args)
+
+
+def _set_tensorboard_writer(args):
+    """Set tensorboard writer."""
+    global _GLOBAL_TENSORBOARD_WRITER
+    _ensure_var_is_not_initialized(_GLOBAL_TENSORBOARD_WRITER,
+                                   'tensorboard writer')
+
+    if hasattr(args, 'tensorboard_dir') and \
+       args.tensorboard_dir and args.rank == (args.world_size - 1):
+        try:
+            from torch.utils.tensorboard import SummaryWriter
+            print('> setting tensorboard ...')
+            _GLOBAL_TENSORBOARD_WRITER = SummaryWriter(
+                log_dir=args.tensorboard_dir,
+                max_queue=args.tensorboard_queue_size)
+        except ModuleNotFoundError:
+            print('WARNING: TensorBoard writing requested but is not '
+                  'available (are you using PyTorch 1.1.0 or later?), '
+                  'no TensorBoard logs will be written.', flush=True)
+
+
+def _set_adlr_autoresume(args):
+    """Initialize ADLR autoresume."""
+    global _GLOBAL_ADLR_AUTORESUME
+    _ensure_var_is_not_initialized(_GLOBAL_ADLR_AUTORESUME, 'adlr autoresume')
+
+    if args.adlr_autoresume:
+        if args.rank == 0:
+            print('enabling autoresume ...', flush=True)
+        sys.path.append(os.environ.get('SUBMIT_SCRIPTS', '.'))
+        try:
+            from userlib.auto_resume import AutoResume
+        except BaseException:
+            print('ADLR autoresume is not available, exiting ...')
+            sys.exit()
+
+        _GLOBAL_ADLR_AUTORESUME = AutoResume
+
+
+def _set_timers():
+    """Initialize timers."""
+    global _GLOBAL_TIMERS
+    _ensure_var_is_not_initialized(_GLOBAL_TIMERS, 'timers')
+    _GLOBAL_TIMERS = Timers()
+
+
+def _ensure_var_is_initialized(var, name):
+    """Make sure the input variable is not None."""
+    assert var is not None, '{} is not initialized.'.format(name)
+
+
+def _ensure_var_is_not_initialized(var, name):
+    """Make sure the input variable is not None."""
+    assert var is None, '{} is already initialized.'.format(name)
+
+
+class _Timer:
+    """Timer."""
+
+    def __init__(self, name):
+        self.name_ = name
+        self.elapsed_ = 0.0
+        self.started_ = False
+        self.start_time = time.time()
+
+    def start(self):
+        """Start the timer."""
+        assert not self.started_, 'timer has already been started'
+        torch.cuda.synchronize()
+        self.start_time = time.time()
+        self.started_ = True
+
+    def stop(self):
+        """Stop the timer."""
+        assert self.started_, 'timer is not started'
+        torch.cuda.synchronize()
+        self.elapsed_ += (time.time() - self.start_time)
+        self.started_ = False
+
+    def reset(self):
+        """Reset timer."""
+        self.elapsed_ = 0.0
+        self.started_ = False
+
+    def elapsed(self, reset=True):
+        """Calculate the elapsed time."""
+        started_ = self.started_
+        # If the timing in progress, end it first.
+        if self.started_:
+            self.stop()
+        # Get the elapsed time.
+        elapsed_ = self.elapsed_
+        # Reset the elapsed time
+        if reset:
+            self.reset()
+        # If timing was in progress, set it back.
+        if started_:
+            self.start()
+        return elapsed_
+
+
+class Timers:
+    """Group of timers."""
+
+    def __init__(self):
+        self.timers = {}
+
+    def __call__(self, name):
+        if name not in self.timers:
+            self.timers[name] = _Timer(name)
+        return self.timers[name]
+
+    def write(self, names, writer, iteration, normalizer=1.0, reset=False):
+        """Write timers to a tensorboard writer"""
+        # currently when using add_scalars,
+        # torch.utils.add_scalars makes each timer its own run, which
+        # polutes the runs list, so we just add each as a scalar
+        assert normalizer > 0.0
+        for name in names:
+            value = self.timers[name].elapsed(reset=reset) / normalizer
+            writer.add_scalar(name + '-time', value, iteration)
+
+    def log(self, names, normalizer=1.0, reset=True):
+        """Log a group of timers."""
+        assert normalizer > 0.0
+        string = 'time (ms)'
+        for name in names:
+            elapsed_time = self.timers[name].elapsed(
+                reset=reset) * 1000.0 / normalizer
+            string += ' | {}: {:.2f}'.format(name, elapsed_time)
+        if torch.distributed.is_initialized():
+            if torch.distributed.get_rank() == (
+                    torch.distributed.get_world_size() - 1):
+                print(string, flush=True)
+        else:
+            print(string, flush=True)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/standalone_gpt.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/standalone_gpt.py
new file mode 100644
index 0000000000..f2720841ee
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/testing/standalone_gpt.py
@@ -0,0 +1,1506 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+
+"""GPT-2 model."""
+import enum
+import math
+
+import torch
+import torch.nn.functional as F
+
+import apex.transformer.utils
+from apex.normalization import FusedLayerNorm as LayerNorm
+from apex.transformer.functional import FusedScaleMaskSoftmax
+from apex.transformer import tensor_parallel
+from apex.transformer import parallel_state
+from apex.transformer.testing.global_vars import get_args
+from apex.transformer.enums import LayerType
+from apex.transformer.enums import AttnType
+from apex.transformer.enums import AttnMaskType
+
+
+_FLOAT_TYPES = (torch.FloatTensor, torch.cuda.FloatTensor)
+_HALF_TYPES = (torch.HalfTensor, torch.cuda.HalfTensor)
+_BF16_TYPES = (torch.BFloat16Tensor, torch.cuda.BFloat16Tensor)
+
+
+class ModelType(enum.Enum):
+    encoder_or_decoder = 1
+    encoder_and_decoder = 2
+
+
+###### BIAS GELU FUSION/ NO AUTOGRAD ################
+# 1/sqrt(2*pi)-> 0.3989423
+# 1/sqrt(2)   -> 0.70710678
+# sqrt(2/pi)  -> 0.79788456
+# this function is tanh approximation of gelu
+# actual gelu is:
+# x * 0.5 * (1.0 + torch.erf(x * 0.70710678))
+
+
+@torch.jit.script
+def bias_gelu(bias, y):
+    x = bias + y
+    return x * 0.5 * (1.0 + torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x)))
+
+
+# gradient of tanh approximation of gelu
+# gradient of actual gelu is:
+# 0.5 * (1. + torch.erf(x * 0.70710678)) + 0.3989423 * x * torch.exp(-0.5 * x * x)
+@torch.jit.script
+def bias_gelu_back(g, bias, y):
+    x = bias + y
+    tanh_out = torch.tanh(0.79788456 * x * (1 + 0.044715 * x * x))
+    # sqrt(2/pi) * 3 * 0.044715 -> 0.1070322243
+    ff = 0.5 * x * ((1 - tanh_out * tanh_out) * (0.79788456 + 0.1070322243 * x * x)) + 0.5 * (1 + tanh_out)
+    return ff * g
+
+
+class MegatronModule(torch.nn.Module):
+    """Megatron specific extensions of torch Module with support
+    for pipelining."""
+
+    def __init__(self, share_word_embeddings=True):
+        super(MegatronModule, self).__init__()
+        self.share_word_embeddings = share_word_embeddings
+
+    def state_dict_for_save_checkpoint(self, destination=None, prefix="", keep_vars=False):
+        """Use this function to override the state dict for
+        saving checkpoints."""
+        return self.state_dict(destination, prefix, keep_vars)
+
+    def word_embeddings_weight(self):
+        if (
+            not parallel_state.is_pipeline_last_stage(ignore_virtual=True)
+            or parallel_state.get_pipeline_model_parallel_world_size() == 1
+        ):
+            return self.language_model.embedding.word_embeddings.weight
+        else:
+            if not self.share_word_embeddings:
+                raise Exception("word_embeddings_weight() called for last " "stage, but share_word_embeddings is false")
+            return self.word_embeddings.weight
+
+    def initialize_word_embeddings(self, init_method_normal):
+        args = get_args()
+        if not self.share_word_embeddings:
+            raise Exception("initialize_word_embeddings() was called but " "share_word_embeddings is false")
+
+        # This function just initializes the word embeddings in the final stage
+        # when we are using pipeline parallelism. Nothing to do if we aren't
+        # using pipeline parallelism.
+        if args.pipeline_model_parallel_size == 1:
+            return
+
+        # Parameters are shared between the word embeddings layers, and the
+        # heads at the end of the model. In a pipelined setup with more than
+        # one stage, the initial embedding layer and the head are on different
+        # workers, so we do the following:
+        # 1. Create a second copy of word_embeddings on the last stage, with
+        #    initial parameters of 0.0.
+        # 2. Do an all-reduce between the first and last stage to ensure that
+        #    the two copies of word_embeddings start off with the same
+        #    parameter values.
+        # 3. In the training loop, before an all-reduce between the grads of
+        #    the two word_embeddings layers to ensure that every applied weight
+        #    update is the same on both stages.
+        if parallel_state.is_pipeline_last_stage():
+            assert not parallel_state.is_pipeline_first_stage()
+            self._word_embeddings_for_head_key = "word_embeddings_for_head"
+            # set word_embeddings weights to 0 here, then copy first
+            # stage's weights using all_reduce below.
+            self.word_embeddings = tensor_parallel.VocabParallelEmbedding(
+                args.padded_vocab_size, args.hidden_size, init_method=init_method_normal(args.init_method_std)
+            )
+            self.word_embeddings.weight.data.fill_(0)
+            self.word_embeddings.weight.shared = True
+
+        # Zero out initial weights for decoder embedding.
+        # NOTE: We don't currently support T5 with the interleaved schedule.
+        if (
+            not parallel_state.is_pipeline_first_stage(ignore_virtual=True)
+            and not parallel_state.is_pipeline_last_stage(ignore_virtual=True)
+            and parallel_state.is_rank_in_embedding_group()
+        ):
+            self.language_model.embedding.zero_parameters()
+
+        # Ensure that first and last stages have the same initial parameter
+        # values.
+        if torch.distributed.is_initialized():
+            if parallel_state.is_rank_in_embedding_group():
+                torch.distributed.all_reduce(
+                    self.word_embeddings_weight().data, group=parallel_state.get_embedding_group()
+                )
+                # All-reduce other embeddings as well as necessary. The last stage
+                # does not have these other embeddings, so just create placeholder
+                # tensors of the right shape with all zeros.
+                # NOTE: We don't currently support T5 with the interleaved schedule.
+                if args.pipeline_model_parallel_split_rank is not None:
+                    # TODO: Support tokentype embedding.
+                    dimensions = (args.max_position_embeddings, args.hidden_size)
+                    if parallel_state.is_pipeline_last_stage(ignore_virtual=True):
+                        position_embeddings = torch.nn.Embedding(*dimensions).cuda()
+                        position_embeddings.weight.data.fill_(0)
+                    else:
+                        self.language_model.embedding.cuda()
+                        position_embeddings = self.language_model.embedding.position_embeddings
+                    torch.distributed.all_reduce(
+                        position_embeddings.weight.data, group=parallel_state.get_embedding_group()
+                    )
+        else:
+            print(
+                "WARNING! Distributed processes aren't initialized, so "
+                "word embeddings in the last layer are not initialized. "
+                "If you are just manipulating a model this is fine, but "
+                "this needs to be handled manually. If you are training "
+                "something is definitely wrong."
+            )
+
+
+class GeLUFunction(torch.autograd.Function):
+    @staticmethod
+    # bias is an optional argument
+    def forward(ctx, input, bias):
+        ctx.save_for_backward(input, bias)
+        return bias_gelu(bias, input)
+
+    @staticmethod
+    def backward(ctx, grad_output):
+        input, bias = ctx.saved_tensors
+        tmp = bias_gelu_back(grad_output, bias, input)
+        return tmp, tmp
+
+
+bias_gelu_impl = GeLUFunction.apply
+
+
+def get_linear_layer(rows, columns, init_method):
+    """Simple linear layer with weight initialization."""
+    layer = torch.nn.Linear(rows, columns)
+    init_method(layer.weight)
+    with torch.no_grad():
+        layer.bias.zero_()
+    return layer
+
+
+def attention_mask_func(attention_scores, attention_mask):
+    attention_scores.masked_fill_(attention_mask, -10000.0)
+    return attention_scores
+
+
+@torch.jit.script
+def gelu_impl(x):
+    """OpenAI's gelu implementation."""
+    return 0.5 * x * (1.0 + torch.tanh(0.7978845608028654 * x * (1.0 + 0.044715 * x * x)))
+
+
+def openai_gelu(x):
+    return gelu_impl(x)
+
+
+# This is actually Python equivalent of torch.nn.functional.gelu(), also with type hints for ONNX exporter
+@torch.jit.script
+def erf_gelu(x):
+    return x * 0.5 * (torch.erf(x / 1.41421).to(dtype=x.dtype) + torch.ones_like(x).to(dtype=x.dtype))
+
+
+def init_method_normal(sigma):
+    """Init method based on N(0, sigma)."""
+
+    def init_(tensor):
+        return torch.nn.init.normal_(tensor, mean=0.0, std=sigma)
+
+    return init_
+
+
+def scaled_init_method_normal(sigma, num_layers):
+    """Init method based on N(0, sigma/sqrt(2*num_layers)."""
+    std = sigma / math.sqrt(2.0 * num_layers)
+
+    def init_(tensor):
+        return torch.nn.init.normal_(tensor, mean=0.0, std=std)
+
+    return init_
+
+
+class ParallelMLP(MegatronModule):
+    """MLP.
+
+    MLP will take the input with h hidden state, project it to 4*h
+    hidden dimension, perform nonlinear transformation, and project the
+    state back into h hidden dimension.
+    """
+
+    def __init__(self, init_method, output_layer_init_method):
+        super().__init__()
+        args = get_args()
+
+        # Project to 4h.
+        self.dense_h_to_4h = tensor_parallel.ColumnParallelLinear(
+            args.hidden_size, args.ffn_hidden_size, gather_output=False, init_method=init_method, skip_bias_add=True
+        )
+
+        self.bias_gelu_fusion = args.bias_gelu_fusion
+        self.activation_func = F.gelu
+        if args.openai_gelu:
+            self.activation_func = openai_gelu
+        elif args.onnx_safe:
+            self.activation_func = erf_gelu
+
+        # Project back to h.
+        self.dense_4h_to_h = tensor_parallel.RowParallelLinear(
+            args.ffn_hidden_size,
+            args.hidden_size,
+            input_is_parallel=True,
+            init_method=output_layer_init_method,
+            skip_bias_add=True,
+        )
+
+    def forward(self, hidden_states):
+
+        # [s, b, 4hp]
+        intermediate_parallel, bias_parallel = self.dense_h_to_4h(hidden_states)
+
+        if self.bias_gelu_fusion:
+            intermediate_parallel = bias_gelu_impl(intermediate_parallel, bias_parallel)
+        else:
+            intermediate_parallel = self.activation_func(intermediate_parallel + bias_parallel)
+
+        # [s, b, h]
+        output, output_bias = self.dense_4h_to_h(intermediate_parallel)
+        return output, output_bias
+
+
+class ParallelAttention(MegatronModule):
+    """Parallel self-attention layer abstract class.
+
+    Self-attention layer takes input with size [b, s, h]
+    and returns output of the same size.
+    """
+
+    def __init__(
+        self,
+        init_method,
+        output_layer_init_method,
+        layer_number,
+        attention_type=AttnType.self_attn,
+        attn_mask_type=AttnMaskType.padding,
+    ):
+        super().__init__()
+        args = get_args()
+        self.fp16 = args.fp16
+        self.bf16 = args.bf16
+
+        self.apply_query_key_layer_scaling = args.apply_query_key_layer_scaling
+        self.attention_softmax_in_fp32 = args.attention_softmax_in_fp32
+        if self.apply_query_key_layer_scaling:
+            self.attention_softmax_in_fp32 = True
+        self.layer_number = max(1, layer_number)
+        self.attention_type = attention_type
+        self.attn_mask_type = attn_mask_type
+        self.params_dtype = args.params_dtype
+
+        projection_size = args.kv_channels * args.num_attention_heads
+
+        # Per attention head and per partition values.
+        world_size = parallel_state.get_tensor_model_parallel_world_size()
+        self.hidden_size_per_partition = apex.transformer.utils.divide(projection_size, world_size)
+        self.hidden_size_per_attention_head = apex.transformer.utils.divide(projection_size, args.num_attention_heads)
+        self.num_attention_heads_per_partition = apex.transformer.utils.divide(args.num_attention_heads, world_size)
+
+        # Strided linear layer.
+        if attention_type == AttnType.self_attn:
+            self.query_key_value = tensor_parallel.ColumnParallelLinear(
+                args.hidden_size, 3 * projection_size, gather_output=False, init_method=init_method
+            )
+        else:
+            assert attention_type == AttnType.cross_attn
+            self.query = tensor_parallel.ColumnParallelLinear(
+                args.hidden_size, projection_size, gather_output=False, init_method=init_method
+            )
+
+            self.key_value = tensor_parallel.ColumnParallelLinear(
+                args.hidden_size, 2 * projection_size, gather_output=False, init_method=init_method
+            )
+
+        coeff = None
+        self.norm_factor = math.sqrt(self.hidden_size_per_attention_head)
+        if self.apply_query_key_layer_scaling:
+            coeff = self.layer_number
+            self.norm_factor *= coeff
+
+        self.scale_mask_softmax = FusedScaleMaskSoftmax(
+            self.fp16,
+            self.bf16,
+            self.attn_mask_type,
+            args.masked_softmax_fusion,
+            attention_mask_func,
+            self.attention_softmax_in_fp32,
+            coeff,
+        )
+
+        # Dropout. Note that for a single iteration, this layer will generate
+        # different outputs on different number of parallel partitions but
+        # on average it should not be partition dependent.
+        self.attention_dropout = torch.nn.Dropout(args.attention_dropout)
+
+        # Output.
+        self.dense = tensor_parallel.RowParallelLinear(
+            projection_size,
+            args.hidden_size,
+            input_is_parallel=True,
+            init_method=output_layer_init_method,
+            skip_bias_add=True,
+        )
+
+        # Inference key-value memory
+        self.inference_key_memory = None
+        self.inference_value_memory = None
+        self.inference_current_sequence_len = 0
+
+    def _allocate_memory(self, inference_max_sequence_len, batch_size):
+        return torch.empty(
+            inference_max_sequence_len,
+            batch_size,
+            self.num_attention_heads_per_partition,
+            self.hidden_size_per_attention_head,
+            dtype=self.params_dtype,
+            device=torch.cuda.current_device(),
+        )
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        encoder_output=None,
+        set_inference_key_value_memory=False,
+        inference_max_sequence_len=None,
+    ):
+        # hidden_states: [sq, b, h]
+
+        # =================================================
+        # Pre-allocate memory for key-values for inference.
+        # =================================================
+        if set_inference_key_value_memory:
+            assert inference_max_sequence_len and inference_max_sequence_len > 0
+            self.inference_key_memory = self._allocate_memory(inference_max_sequence_len, hidden_states.size(1))
+            self.inference_value_memory = self._allocate_memory(inference_max_sequence_len, hidden_states.size(1))
+            self.inference_current_sequence_len = 0
+        # Some consistency check.
+        if inference_max_sequence_len:
+            assert self.inference_current_sequence_len < self.inference_key_memory.size(0)
+            assert inference_max_sequence_len == self.inference_key_memory.size(0)
+        # This is added for safety. In case inference_max_sequence_len
+        # is not provided, make sure there is no potential memory left
+        # from previous inference.
+        if not inference_max_sequence_len:
+            self.inference_key_memory = None
+            self.inference_value_memory = None
+
+        # =====================
+        # Query, Key, and Value
+        # =====================
+
+        if self.attention_type == AttnType.self_attn:
+            # Attention heads [sq, b, h] --> [sq, b, (np * 3 * hn)]
+            mixed_x_layer, _ = self.query_key_value(hidden_states)
+
+            # [sq, b, (np * 3 * hn)] --> [sq, b, np, 3 * hn]
+            new_tensor_shape = mixed_x_layer.size()[:-1] + (
+                self.num_attention_heads_per_partition,
+                3 * self.hidden_size_per_attention_head,
+            )
+            mixed_x_layer = mixed_x_layer.view(*new_tensor_shape)
+
+            # [sq, b, np, 3 * hn] --> 3 [sq, b, np, hn]
+            (query_layer, key_layer, value_layer) = tensor_parallel.split_tensor_along_last_dim(mixed_x_layer, 3)
+        else:
+            # Attention heads [sk, b, h] --> [sk, b, (np * 2 * hn)]
+            mixed_kv_layer, _ = self.key_value(encoder_output)
+
+            # [sk, b, (np * 2 * hn)] --> [sk, b, np, 2 * hn]
+            new_tensor_shape = mixed_kv_layer.size()[:-1] + (
+                self.num_attention_heads_per_partition,
+                2 * self.hidden_size_per_attention_head,
+            )
+            mixed_kv_layer = mixed_kv_layer.view(*new_tensor_shape)
+
+            # [sk, b, np, 2 * hn] --> 2 [sk, b, np, hn]
+            (key_layer, value_layer) = tensor_parallel.split_tensor_along_last_dim(mixed_kv_layer, 2)
+
+            # Attention head [sq, b, h] --> [sq, b, hp]
+            query_layer, _ = self.query(hidden_states)
+            # [sq, b, hp] --> [sq, b, np, hn]
+            new_tensor_shape = query_layer.size()[:-1] + (
+                self.num_attention_heads_per_partition,
+                self.hidden_size_per_attention_head,
+            )
+            query_layer = query_layer.view(*new_tensor_shape)
+
+        # ===================================================
+        # Adjust key, value, and attention mask for inference
+        # ===================================================
+
+        if inference_max_sequence_len:
+            # Adjust the range variables.
+            start = self.inference_current_sequence_len
+            self.inference_current_sequence_len += key_layer.size(0)
+            end = self.inference_current_sequence_len
+            # Copy key and values.
+            self.inference_key_memory[start:end, ...] = key_layer
+            self.inference_value_memory[start:end, ...] = value_layer
+            key_layer = self.inference_key_memory[:end, ...]
+            value_layer = self.inference_value_memory[:end, ...]
+            # Adjust attention mask
+            attention_mask = attention_mask[..., start:end, :end]
+
+        # ===================================
+        # Raw attention scores. [b, np, s, s]
+        # ===================================
+
+        # [b, np, sq, sk]
+        output_size = (query_layer.size(1), query_layer.size(2), query_layer.size(0), key_layer.size(0))
+
+        # [sq, b, np, hn] -> [sq, b * np, hn]
+        query_layer = query_layer.view(output_size[2], output_size[0] * output_size[1], -1)
+        # [sk, b, np, hn] -> [sk, b * np, hn]
+        key_layer = key_layer.view(output_size[3], output_size[0] * output_size[1], -1)
+
+        # preallocting result tensor: [b * np, sq, sk]
+        matmul_result = torch.empty(
+            output_size[0] * output_size[1],
+            output_size[2],
+            output_size[3],
+            dtype=query_layer.dtype,
+            device=torch.cuda.current_device(),
+        )
+
+        # Raw attention scores. [b * np, sq, sk]
+        matmul_result = torch.baddbmm(
+            matmul_result,
+            query_layer.transpose(0, 1),  # [b * np, sq, hn]
+            key_layer.transpose(0, 1).transpose(1, 2),  # [b * np, hn, sk]
+            beta=0.0,
+            alpha=(1.0 / self.norm_factor),
+        )
+
+        # change view to [b, np, sq, sk]
+        attention_scores = matmul_result.view(*output_size)
+
+        # ===========================
+        # Attention probs and dropout
+        # ===========================
+
+        # attention scores and attention mask [b, np, sq, sk]
+        attention_probs = self.scale_mask_softmax(attention_scores, attention_mask)
+
+        # This is actually dropping out entire tokens to attend to, which might
+        # seem a bit unusual, but is taken from the original Transformer paper.
+        with tensor_parallel.get_cuda_rng_tracker().fork():
+            attention_probs = self.attention_dropout(attention_probs)
+
+        # =========================
+        # Context layer. [sq, b, hp]
+        # =========================
+
+        # value_layer -> context layer.
+        # [sk, b, np, hn] --> [b, np, sq, hn]
+
+        # context layer shape: [b, np, sq, hn]
+        output_size = (value_layer.size(1), value_layer.size(2), query_layer.size(0), value_layer.size(3))
+
+        # change view [sk, b * np, hn]
+        value_layer = value_layer.view(value_layer.size(0), output_size[0] * output_size[1], -1)
+
+        # change view [b * np, sq, sk]
+        attention_probs = attention_probs.view(output_size[0] * output_size[1], output_size[2], -1)
+
+        # matmul: [b * np, sq, hn]
+        context_layer = torch.bmm(attention_probs, value_layer.transpose(0, 1))
+
+        # change view [b, np, sq, hn]
+        context_layer = context_layer.view(*output_size)
+
+        # [b, np, sq, hn] --> [sq, b, np, hn]
+        context_layer = context_layer.permute(2, 0, 1, 3).contiguous()
+
+        # [sq, b, np, hn] --> [sq, b, hp]
+        new_context_layer_shape = context_layer.size()[:-2] + (self.hidden_size_per_partition,)
+        context_layer = context_layer.view(*new_context_layer_shape)
+
+        # =================
+        # Output. [sq, b, h]
+        # =================
+
+        output, bias = self.dense(context_layer)
+
+        return output, bias
+
+
+@torch.jit.script
+def bias_dropout_add(x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float, training: bool) -> torch.Tensor:
+    out = torch.nn.functional.dropout(x + bias, p=prob, training=training)
+    out = residual + out
+    return out
+
+
+def get_bias_dropout_add(training):
+    def _bias_dropout_add(x, bias, residual, prob):
+        return bias_dropout_add(x, bias, residual, prob, training)
+
+    return _bias_dropout_add
+
+
+@torch.jit.script
+def bias_dropout_add_fused_train(
+    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
+) -> torch.Tensor:
+    return bias_dropout_add(x, bias, residual, prob, True)
+
+
+@torch.jit.script
+def bias_dropout_add_fused_inference(
+    x: torch.Tensor, bias: torch.Tensor, residual: torch.Tensor, prob: float
+) -> torch.Tensor:
+    return bias_dropout_add(x, bias, residual, prob, False)
+
+
+class ParallelTransformerLayer(MegatronModule):
+    """A single transformer layer.
+
+    Transformer layer takes input with size [b, s, h] and returns an
+    output of the same size.
+    """
+
+    def __init__(
+        self,
+        init_method,
+        output_layer_init_method,
+        layer_number,
+        layer_type=LayerType.encoder,
+        self_attn_mask_type=AttnMaskType.padding,
+    ):
+        args = get_args()
+
+        super().__init__()
+        self.layer_number = layer_number
+        self.layer_type = layer_type
+
+        self.apply_residual_connection_post_layernorm = args.apply_residual_connection_post_layernorm
+
+        self.bf16 = args.bf16
+        self.fp32_residual_connection = args.fp32_residual_connection
+
+        # Layernorm on the input data.
+        self.input_layernorm = LayerNorm(args.hidden_size, eps=args.layernorm_epsilon)
+
+        # Self attention.
+        self.self_attention = ParallelAttention(
+            init_method,
+            output_layer_init_method,
+            layer_number,
+            attention_type=AttnType.self_attn,
+            attn_mask_type=self_attn_mask_type,
+        )
+        self.hidden_dropout = args.hidden_dropout
+        self.bias_dropout_fusion = args.bias_dropout_fusion
+
+        # Layernorm on the attention output
+        self.post_attention_layernorm = LayerNorm(args.hidden_size, eps=args.layernorm_epsilon)
+
+        if self.layer_type == LayerType.decoder:
+            self.inter_attention = ParallelAttention(
+                init_method, output_layer_init_method, layer_number, attention_type=AttnType.cross_attn
+            )
+            # Layernorm on the attention output.
+            self.post_inter_attention_layernorm = LayerNorm(args.hidden_size, eps=args.layernorm_epsilon)
+
+        # MLP
+        self.mlp = ParallelMLP(init_method, output_layer_init_method)
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        encoder_output=None,
+        enc_dec_attn_mask=None,
+        set_inference_key_value_memory=False,
+        inference_max_sequence_len=None,
+    ):
+        # hidden_states: [b, s, h]
+
+        # Layer norm at the beginning of the transformer layer.
+        layernorm_output = self.input_layernorm(hidden_states)
+        # Self attention.
+        attention_output, attention_bias = self.self_attention(
+            layernorm_output,
+            attention_mask,
+            set_inference_key_value_memory=set_inference_key_value_memory,
+            inference_max_sequence_len=inference_max_sequence_len,
+        )
+
+        # Residual connection.
+        if self.apply_residual_connection_post_layernorm:
+            residual = layernorm_output
+        else:
+            residual = hidden_states
+
+        # jit scripting for a nn.module (with dropout) is not
+        # trigerring the fusion kernel. For now, we use two
+        # different nn.functional routines to account for varying
+        # dropout semantics during training and inference phases.
+        if self.bias_dropout_fusion:
+            if self.training:
+                bias_dropout_add_func = bias_dropout_add_fused_train
+            else:
+                bias_dropout_add_func = bias_dropout_add_fused_inference
+        else:
+            bias_dropout_add_func = get_bias_dropout_add(self.training)
+
+        # re-enable torch grad to enable fused optimization.
+        with torch.enable_grad():
+            layernorm_input = bias_dropout_add_func(
+                attention_output, attention_bias.expand_as(residual), residual, self.hidden_dropout
+            )
+
+        # Layer norm post the self attention.
+        layernorm_output = self.post_attention_layernorm(layernorm_input)
+
+        if self.layer_type == LayerType.decoder:
+            attention_output, attention_bias = self.inter_attention(
+                layernorm_output, enc_dec_attn_mask, encoder_output=encoder_output
+            )
+            # residual connection
+            if self.apply_residual_connection_post_layernorm:
+                residual = layernorm_output
+            else:
+                residual = layernorm_input
+
+            # re-enable torch grad to enable fused optimization.
+            with torch.enable_grad():
+                layernorm_input = bias_dropout_add_func(
+                    attention_output, attention_bias.expand_as(residual), residual, self.hidden_dropout
+                )
+
+            # Layer norm post the decoder attention
+            layernorm_output = self.post_inter_attention_layernorm(layernorm_input)
+
+        # MLP.
+        mlp_output, mlp_bias = self.mlp(layernorm_output)
+
+        # Second residual connection.
+        if self.apply_residual_connection_post_layernorm:
+            residual = layernorm_output
+        else:
+            residual = layernorm_input
+
+        # re-enable torch grad to enable fused optimization.
+        with torch.enable_grad():
+            output = bias_dropout_add_func(mlp_output, mlp_bias.expand_as(residual), residual, self.hidden_dropout)
+
+        return output
+
+
+class ParallelTransformer(MegatronModule):
+    """Transformer class."""
+
+    def __init__(
+        self,
+        init_method,
+        output_layer_init_method,
+        layer_type=LayerType.encoder,
+        self_attn_mask_type=AttnMaskType.padding,
+        pre_process=True,
+        post_process=True,
+    ):
+        super().__init__()
+        args = get_args()
+
+        self.bf16 = args.bf16
+        self.fp32_residual_connection = args.fp32_residual_connection
+        self.pre_process = pre_process
+        self.post_process = post_process
+        self.input_tensor = None
+
+        # Store activation checkpointing flag.
+        self.activations_checkpoint_method = args.activations_checkpoint_method
+        self.activations_checkpoint_num_layers = args.activations_checkpoint_num_layers
+        self.distribute_checkpointed_activations = args.distribute_checkpointed_activations
+
+        num_layers = args.num_layers
+        # Number of layers.
+        assert (
+            num_layers % parallel_state.get_pipeline_model_parallel_world_size() == 0
+        ), "num_layers must be divisible by pipeline_model_parallel_size"
+        self.num_layers = num_layers // parallel_state.get_pipeline_model_parallel_world_size()
+
+        # Transformer layers.
+        def build_layer(layer_number):
+            return ParallelTransformerLayer(
+                init_method,
+                output_layer_init_method,
+                layer_number,
+                layer_type=layer_type,
+                self_attn_mask_type=self_attn_mask_type,
+            )
+
+        if args.virtual_pipeline_model_parallel_size is not None:
+            assert args.num_layers % args.virtual_pipeline_model_parallel_size == 0, (
+                "num_layers_per_stage must be divisible by " "virtual_pipeline_model_parallel_size"
+            )
+            # Number of layers in each model chunk is the number of layers in the stage,
+            # divided by the number of model chunks in a stage.
+            self.num_layers = self.num_layers // args.virtual_pipeline_model_parallel_size
+            # With 8 layers, 2 stages, and 4 model chunks, we want an assignment of
+            # layers to stages like (each list is a model chunk):
+            # Stage 0: [0]  [2]  [4]  [6]
+            # Stage 1: [1]  [3]  [5]  [7]
+            # With 8 layers, 2 stages, and 2 virtual stages, we want an assignment of
+            # layers to stages like (each list is a model chunk):
+            # Stage 0: [0, 1]  [4, 5]
+            # Stage 1: [2, 3]  [6, 7]
+            offset = parallel_state.get_virtual_pipeline_model_parallel_rank() * (
+                args.num_layers // args.virtual_pipeline_model_parallel_size
+            ) + (parallel_state.get_pipeline_model_parallel_rank() * self.num_layers)
+        else:
+            # Each stage gets a contiguous set of layers.
+            offset = parallel_state.get_pipeline_model_parallel_rank() * self.num_layers
+
+        self.layers = torch.nn.ModuleList([build_layer(i + 1 + offset) for i in range(self.num_layers)])
+
+        if self.post_process:
+            # Final layer norm before output.
+            self.final_layernorm = LayerNorm(args.hidden_size, eps=args.layernorm_epsilon)
+
+    def _get_layer(self, layer_number):
+        return self.layers[layer_number]
+
+    def _checkpointed_forward(self, hidden_states, attention_mask, encoder_output, enc_dec_attn_mask):
+        """Forward method with activation checkpointing."""
+
+        def custom(start, end):
+            def custom_forward(*inputs):
+                x_ = inputs[0]
+                attention_mask = inputs[1]
+                encoder_output = inputs[2]
+                enc_dec_attn_mask = inputs[3]
+                for index in range(start, end):
+                    layer = self._get_layer(index)
+                    x_ = layer(x_, attention_mask, encoder_output, enc_dec_attn_mask)
+                return x_
+
+            return custom_forward
+
+        def distribute_checkpointed_activations_helper(layer_number):
+            """Distribute checkpointed activations across the tensor model
+               Parallel ranks if the `distribute-checkpointed-activations
+               is on and either of the following conditions is met:
+                 - it is not the first layer in the in the pipeline stage.
+                   The first layer is used in the pipeline parallelism
+                   and changing its shape throws error in the backward pass.
+                 - we are at the first pipline stage so the input tensor is
+                   not used in pipeline parallelism. Note that no pipeline
+                   parallelism is a special case of this.
+            """
+            not_first_layer_in_pipeline_stage = layer_number > 0
+            is_first_pipeline_stage = parallel_state.get_pipeline_model_parallel_rank() == 0
+            return self.distribute_checkpointed_activations and (
+                not_first_layer_in_pipeline_stage or is_first_pipeline_stage
+            )
+
+        if self.activations_checkpoint_method == "uniform":
+            # Uniformly divide the total number of Transformer layers and checkpoint
+            # the input activation of each divided chunk.
+            # A method to further reduce memory usage reducing checkpoints.
+            l = 0
+            while l < self.num_layers:
+                hidden_states = tensor_parallel.checkpoint(
+                    custom(l, l + self.activations_checkpoint_num_layers),
+                    distribute_checkpointed_activations_helper(l),
+                    hidden_states,
+                    attention_mask,
+                    encoder_output,
+                    enc_dec_attn_mask,
+                )
+                l += self.activations_checkpoint_num_layers
+        elif self.activations_checkpoint_method == "block":
+            # Checkpoint the input activation of only a set number of individual
+            # Transformer layers and skip the rest.
+            # A method fully use the device memory removing redundant re-computation.
+            for l in range(self.num_layers):
+                if l < self.activations_checkpoint_num_layers:
+                    hidden_states = tensor_parallel.checkpoint(
+                        custom(l, l + 1),
+                        distribute_checkpointed_activations_helper(l),
+                        hidden_states,
+                        attention_mask,
+                        encoder_output,
+                        enc_dec_attn_mask,
+                    )
+                else:
+                    hidden_states = custom(l, l + 1)(hidden_states, attention_mask, encoder_output, enc_dec_attn_mask)
+        else:
+            raise ValueError("Invalid activation checkpoint method.")
+
+        return hidden_states
+
+    def set_input_tensor(self, input_tensor):
+        """Set input tensor to be used instead of forward()'s input.
+
+        When doing pipeline parallelism the input from the previous
+        stage comes from communication, not from the input, so the
+        model's forward_step_func won't have it. This function is thus
+        used by internal code to bypass the input provided by the
+        forward_step_func"""
+        self.input_tensor = input_tensor
+
+    def forward(
+        self,
+        hidden_states,
+        attention_mask,
+        encoder_output=None,
+        enc_dec_attn_mask=None,
+        set_inference_key_value_memory=False,
+        inference_max_sequence_len=None,
+    ):
+
+        # Checks.
+        if inference_max_sequence_len:
+            assert self.activations_checkpoint_method is None, "inference does not work with activation checkpointing"
+
+        if self.pre_process:
+            # Data format change to avoid explicit tranposes : [b s h] --> [s b h].
+            # If the input flag for fp32 residual connection is set, convert for float.
+            if self.fp32_residual_connection:
+                hidden_states = hidden_states.transpose(0, 1).contiguous().float()
+            # Otherwise, leave it as is.
+            else:
+                hidden_states = hidden_states.transpose(0, 1).contiguous()
+        else:
+            # See set_input_tensor()
+            hidden_states = self.input_tensor
+
+        if encoder_output is not None:
+            encoder_output = encoder_output.transpose(0, 1).contiguous()
+
+        if self.activations_checkpoint_method is not None:
+            hidden_states = self._checkpointed_forward(hidden_states, attention_mask, encoder_output, enc_dec_attn_mask)
+        else:
+            for index in range(self.num_layers):
+                layer = self._get_layer(index)
+                hidden_states = layer(
+                    hidden_states,
+                    attention_mask,
+                    encoder_output=encoder_output,
+                    enc_dec_attn_mask=enc_dec_attn_mask,
+                    set_inference_key_value_memory=set_inference_key_value_memory,
+                    inference_max_sequence_len=inference_max_sequence_len,
+                )
+
+        # Final layer norm.
+        if self.post_process:
+            # Reverting data format change [s b h] --> [b s h].
+            hidden_states = hidden_states.transpose(0, 1).contiguous()
+            output = self.final_layernorm(hidden_states)
+        else:
+            output = hidden_states
+
+        return output
+
+
+def parallel_lm_logits(input_, word_embeddings_weight, parallel_output, bias=None):
+    """LM logits using word embedding weights."""
+    # Parallel logits.
+    input_parallel = tensor_parallel.copy_to_tensor_model_parallel_region(input_)
+    # Matrix multiply.
+    if bias is None:
+        logits_parallel = F.linear(input_parallel, word_embeddings_weight)
+    else:
+        logits_parallel = F.linear(input_parallel, word_embeddings_weight, bias)
+    # Gather if needed.
+    if parallel_output:
+        return logits_parallel
+
+    return tensor_parallel.gather_from_tensor_model_parallel_region(logits_parallel)
+
+
+def get_language_model(
+    num_tokentypes,
+    add_pooler,
+    encoder_attn_mask_type,
+    init_method=None,
+    scaled_init_method=None,
+    add_encoder=True,
+    add_decoder=False,
+    decoder_attn_mask_type=AttnMaskType.causal,
+    pre_process=True,
+    post_process=True,
+):
+    """Build language model and return along with the key to save."""
+    args = get_args()
+
+    if init_method is None:
+        init_method = init_method_normal(args.init_method_std)
+    if scaled_init_method is None:
+        scaled_init_method = scaled_init_method_normal(args.init_method_std, args.num_layers)
+
+    # Language model.
+    language_model = TransformerLanguageModel(
+        init_method,
+        scaled_init_method,
+        encoder_attn_mask_type,
+        num_tokentypes=num_tokentypes,
+        add_encoder=add_encoder,
+        add_decoder=add_decoder,
+        decoder_attn_mask_type=decoder_attn_mask_type,
+        add_pooler=add_pooler,
+        pre_process=pre_process,
+        post_process=post_process,
+    )
+    # key used for checkpoints.
+    language_model_key = "language_model"
+
+    return language_model, language_model_key
+
+
+class Pooler(MegatronModule):
+    """Pooler layer.
+
+    Pool hidden states of a specific token (for example start of the
+    sequence) and add a linear transformation followed by a tanh.
+
+    Arguments:
+        hidden_size: hidden size
+        init_method: weight initialization method for the linear layer.
+            bias is set to zero.
+    """
+
+    def __init__(self, hidden_size, init_method):
+        super(Pooler, self).__init__()
+        self.dense = get_linear_layer(hidden_size, hidden_size, init_method)
+
+    def forward(self, hidden_states, sequence_index=0):
+        # hidden_states: [b, s, h]
+        # sequence_index: index of the token to pool.
+        pooled = hidden_states[:, sequence_index, :]
+        pooled = self.dense(pooled)
+        pooled = torch.tanh(pooled)
+        return pooled
+
+
+class Embedding(MegatronModule):
+    """Language model embeddings.
+
+    Arguments:
+        hidden_size: hidden size
+        vocab_size: vocabulary size
+        max_sequence_length: maximum size of sequence. This
+                             is used for positional embedding
+        embedding_dropout_prob: dropout probability for embeddings
+        init_method: weight initialization method
+        num_tokentypes: size of the token-type embeddings. 0 value
+                        will ignore this embedding
+    """
+
+    def __init__(
+        self, hidden_size, vocab_size, max_sequence_length, embedding_dropout_prob, init_method, num_tokentypes=0
+    ):
+        super().__init__()
+
+        self.hidden_size = hidden_size
+        self.init_method = init_method
+        self.num_tokentypes = num_tokentypes
+
+        # Word embeddings (parallel).
+        self.word_embeddings = tensor_parallel.VocabParallelEmbedding(
+            vocab_size, self.hidden_size, init_method=self.init_method
+        )
+        self._word_embeddings_key = "word_embeddings"
+
+        # Position embedding (serial).
+        self.position_embeddings = torch.nn.Embedding(max_sequence_length, self.hidden_size)
+        self._position_embeddings_key = "position_embeddings"
+        # Initialize the position embeddings.
+        self.init_method(self.position_embeddings.weight)
+
+        # Token type embedding.
+        # Add this as an optional field that can be added through
+        # method call so we can load a pretrain model without
+        # token types and add them as needed.
+        self._tokentype_embeddings_key = "tokentype_embeddings"
+        if self.num_tokentypes > 0:
+            self.tokentype_embeddings = torch.nn.Embedding(self.num_tokentypes, self.hidden_size)
+            # Initialize the token-type embeddings.
+            self.init_method(self.tokentype_embeddings.weight)
+        else:
+            self.tokentype_embeddings = None
+
+        # Embeddings dropout
+        self.embedding_dropout = torch.nn.Dropout(embedding_dropout_prob)
+
+    def zero_parameters(self):
+        """Zero out all parameters in embedding."""
+        self.word_embeddings.weight.data.fill_(0)
+        self.word_embeddings.weight.shared = True
+        self.position_embeddings.weight.data.fill_(0)
+        self.position_embeddings.weight.shared = True
+        if self.num_tokentypes > 0:
+            self.tokentype_embeddings.weight.data.fill_(0)
+            self.tokentype_embeddings.weight.shared = True
+
+    def add_tokentype_embeddings(self, num_tokentypes):
+        """Add token-type embedding. This function is provided so we can add
+        token-type embeddings in case the pretrained model does not have it.
+        This allows us to load the model normally and then add this embedding.
+        """
+        if self.tokentype_embeddings is not None:
+            raise Exception("tokentype embeddings is already initialized")
+        if torch.distributed.get_rank() == 0:
+            print("adding embedding for {} tokentypes".format(num_tokentypes), flush=True)
+        self.num_tokentypes = num_tokentypes
+        self.tokentype_embeddings = torch.nn.Embedding(num_tokentypes, self.hidden_size)
+        # Initialize the token-type embeddings.
+        self.init_method(self.tokentype_embeddings.weight)
+
+    def forward(self, input_ids, position_ids, tokentype_ids=None):
+        # Embeddings.
+        words_embeddings = self.word_embeddings(input_ids)
+        position_embeddings = self.position_embeddings(position_ids)
+        embeddings = words_embeddings + position_embeddings
+        if tokentype_ids is not None:
+            assert self.tokentype_embeddings is not None
+            embeddings = embeddings + self.tokentype_embeddings(tokentype_ids)
+        else:
+            assert self.tokentype_embeddings is None
+
+        # Dropout.
+        embeddings = self.embedding_dropout(embeddings)
+
+        return embeddings
+
+    def state_dict_for_save_checkpoint(self, destination=None, prefix="", keep_vars=False):
+        """For easy load."""
+
+        state_dict_ = {}
+        state_dict_[self._word_embeddings_key] = self.word_embeddings.state_dict(destination, prefix, keep_vars)
+        state_dict_[self._position_embeddings_key] = self.position_embeddings.state_dict(destination, prefix, keep_vars)
+        if self.num_tokentypes > 0:
+            state_dict_[self._tokentype_embeddings_key] = self.tokentype_embeddings.state_dict(
+                destination, prefix, keep_vars
+            )
+
+        return state_dict_
+
+    def load_state_dict(self, state_dict, strict=True):
+        """Customized load."""
+
+        # Word embedding.
+        if self._word_embeddings_key in state_dict:
+            state_dict_ = state_dict[self._word_embeddings_key]
+        else:
+            # for backward compatibility.
+            state_dict_ = {}
+            for key in state_dict.keys():
+                if "word_embeddings" in key:
+                    state_dict_[key.split("word_embeddings.")[1]] = state_dict[key]
+        self.word_embeddings.load_state_dict(state_dict_, strict=strict)
+
+        # Position embedding.
+        if self._position_embeddings_key in state_dict:
+            state_dict_ = state_dict[self._position_embeddings_key]
+        else:
+            # for backward compatibility.
+            state_dict_ = {}
+            for key in state_dict.keys():
+                if "position_embeddings" in key:
+                    state_dict_[key.split("position_embeddings.")[1]] = state_dict[key]
+        self.position_embeddings.load_state_dict(state_dict_, strict=strict)
+
+        # Tokentype embedding.
+        if self.num_tokentypes > 0:
+            state_dict_ = {}
+            if self._tokentype_embeddings_key in state_dict:
+                state_dict_ = state_dict[self._tokentype_embeddings_key]
+            else:
+                # for backward compatibility.
+                for key in state_dict.keys():
+                    if "tokentype_embeddings" in key:
+                        state_dict_[key.split("tokentype_embeddings.")[1]] = state_dict[key]
+            if len(state_dict_.keys()) > 0:
+                self.tokentype_embeddings.load_state_dict(state_dict_, strict=strict)
+            else:
+                print(
+                    "***WARNING*** expected tokentype embeddings in the " "checkpoint but could not find it", flush=True
+                )
+
+
+class TransformerLanguageModel(MegatronModule):
+    """Transformer language model.
+
+    Arguments:
+        transformer_hparams: transformer hyperparameters
+        vocab_size: vocabulary size
+        max_sequence_length: maximum size of sequence. This
+                             is used for positional embedding
+        embedding_dropout_prob: dropout probability for embeddings
+        num_tokentypes: size of the token-type embeddings. 0 value
+                        will ignore this embedding
+    """
+
+    def __init__(
+        self,
+        init_method,
+        output_layer_init_method,
+        encoder_attn_mask_type,
+        num_tokentypes=0,
+        add_encoder=True,
+        add_decoder=False,
+        decoder_attn_mask_type=AttnMaskType.causal,
+        add_pooler=False,
+        pre_process=True,
+        post_process=True,
+    ):
+        super().__init__()
+        args = get_args()
+
+        self.pre_process = pre_process
+        self.post_process = post_process
+        self.hidden_size = args.hidden_size
+        self.num_tokentypes = num_tokentypes
+        self.init_method = init_method
+        self.add_encoder = add_encoder
+        self.encoder_attn_mask_type = encoder_attn_mask_type
+        self.add_decoder = add_decoder
+        self.decoder_attn_mask_type = decoder_attn_mask_type
+        self.add_pooler = add_pooler
+        self.encoder_hidden_state = None
+
+        # Embeddings.
+        if self.pre_process:
+            self.embedding = Embedding(
+                self.hidden_size,
+                args.padded_vocab_size,
+                args.max_position_embeddings,
+                args.hidden_dropout,
+                self.init_method,
+                self.num_tokentypes,
+            )
+            self._embedding_key = "embedding"
+
+        # Transformer.
+        # Encoder (usually set to True, False if part of an encoder-decoder
+        # architecture and in encoder-only stage).
+        if self.add_encoder:
+            self.encoder = ParallelTransformer(
+                self.init_method,
+                output_layer_init_method,
+                self_attn_mask_type=self.encoder_attn_mask_type,
+                pre_process=self.pre_process,
+                post_process=self.post_process,
+            )
+            self._encoder_key = "encoder"
+        else:
+            self.encoder = None
+
+        # Decoder (usually set to False, True if part of an encoder-decoder
+        # architecture and in decoder-only stage).
+        if self.add_decoder:
+            # Temporary assertion until we verify correctness of pipeline parallelism
+            # implementation of T5.
+            assert (
+                args.pipeline_model_parallel_size == 1
+            ), "pipeline parallelism is not supported in the presence of decoder"
+            self.decoder = ParallelTransformer(
+                self.init_method,
+                output_layer_init_method,
+                layer_type=LayerType.decoder,
+                self_attn_mask_type=self.decoder_attn_mask_type,
+                pre_process=self.pre_process,
+                post_process=self.post_process,
+            )
+            self._decoder_key = "decoder"
+        else:
+            self.decoder = None
+
+        if self.post_process:
+            # Pooler.
+            if self.add_pooler:
+                self.pooler = Pooler(self.hidden_size, self.init_method)
+                self._pooler_key = "pooler"
+
+    def set_input_tensor(self, input_tensor):
+        """ See megatron.model.transformer.set_input_tensor()"""
+
+        # This is usually handled in schedules.py but some inference code still
+        # gives us non-lists or None
+        if not isinstance(input_tensor, list):
+            input_tensor = [input_tensor]
+
+        if self.add_encoder and self.add_decoder:
+            assert (
+                len(input_tensor) == 1
+            ), "input_tensor should only be length 1 for stage with both encoder and decoder"
+            self.encoder.set_input_tensor(input_tensor[0])
+        elif self.add_encoder:
+            assert len(input_tensor) == 1, "input_tensor should only be length 1 for stage with only encoder"
+            self.encoder.set_input_tensor(input_tensor[0])
+        elif self.add_decoder:
+            if len(input_tensor) == 2:
+                self.decoder.set_input_tensor(input_tensor[0])
+                self.encoder_hidden_state = input_tensor[1]
+            elif len(input_tensor) == 1:
+                self.decoder.set_input_tensor(None)
+                self.encoder_hidden_state = input_tensor[0]
+            else:
+                raise Exception("input_tensor must have either length 1 or 2")
+        else:
+            raise Exception("Stage must have at least either encoder or decoder")
+
+    def forward(
+        self,
+        enc_input_ids,
+        enc_position_ids,
+        enc_attn_mask,
+        dec_input_ids=None,
+        dec_position_ids=None,
+        dec_attn_mask=None,
+        enc_dec_attn_mask=None,
+        tokentype_ids=None,
+        set_inference_key_value_memory=False,
+        inference_max_sequence_len=None,
+        pooling_sequence_index=0,
+        enc_hidden_states=None,
+        output_enc_hidden=False,
+    ):
+
+        # Encoder embedding.
+        if self.pre_process:
+            encoder_input = self.embedding(enc_input_ids, enc_position_ids, tokentype_ids=tokentype_ids)
+        else:
+            encoder_input = None
+
+        # Run encoder.
+        if enc_hidden_states is None:
+            if self.encoder is not None:
+                encoder_output = self.encoder(
+                    encoder_input,
+                    enc_attn_mask,
+                    set_inference_key_value_memory=set_inference_key_value_memory,
+                    inference_max_sequence_len=inference_max_sequence_len,
+                )
+            else:
+                encoder_output = self.encoder_hidden_state
+        else:
+            encoder_output = enc_hidden_states.to(encoder_input.dtype)
+
+        if self.post_process:
+            if self.add_pooler:
+                pooled_output = self.pooler(encoder_output, pooling_sequence_index)
+
+        # output_enc_hidden refers to when we just need the encoder's
+        # output. For example, it is helpful to compute
+        # similarity between two sequences by average pooling
+        if not self.add_decoder or output_enc_hidden:
+            if self.add_pooler and self.post_process:
+                return encoder_output, pooled_output
+            else:
+                return encoder_output
+
+        # Decoder embedding.
+        if self.pre_process:
+            decoder_input = self.embedding(dec_input_ids, dec_position_ids)
+        else:
+            decoder_input = None
+
+        # Run decoder.
+        decoder_output = self.decoder(
+            decoder_input,
+            dec_attn_mask,
+            encoder_output=encoder_output,
+            enc_dec_attn_mask=enc_dec_attn_mask,
+            set_inference_key_value_memory=set_inference_key_value_memory,
+            inference_max_sequence_len=inference_max_sequence_len,
+        )
+
+        if self.add_pooler and self.post_process:
+            return decoder_output, encoder_output, pooled_output
+        else:
+            return decoder_output, encoder_output
+
+    def state_dict_for_save_checkpoint(self, destination=None, prefix="", keep_vars=False):
+        """For easy load."""
+
+        state_dict_ = {}
+        if self.pre_process:
+            state_dict_[self._embedding_key] = self.embedding.state_dict_for_save_checkpoint(
+                destination, prefix, keep_vars
+            )
+        if self.add_encoder:
+            state_dict_[self._encoder_key] = self.encoder.state_dict_for_save_checkpoint(destination, prefix, keep_vars)
+        if self.post_process:
+            if self.add_pooler:
+                state_dict_[self._pooler_key] = self.pooler.state_dict_for_save_checkpoint(
+                    destination, prefix, keep_vars
+                )
+        if self.add_decoder:
+            state_dict_[self._decoder_key] = self.decoder.state_dict_for_save_checkpoint(destination, prefix, keep_vars)
+
+        return state_dict_
+
+    def load_state_dict(self, state_dict, strict=True):
+        """Customized load."""
+
+        # Embedding.
+        if self.pre_process:
+            if self._embedding_key in state_dict:
+                state_dict_ = state_dict[self._embedding_key]
+            else:
+                # for backward compatibility.
+                state_dict_ = {}
+                for key in state_dict.keys():
+                    if "_embeddings" in key:
+                        state_dict_[key] = state_dict[key]
+            self.embedding.load_state_dict(state_dict_, strict=strict)
+
+        # Encoder.
+        if self.add_encoder:
+            if self._encoder_key in state_dict:
+                state_dict_ = state_dict[self._encoder_key]
+            # For backward compatibility.
+            elif "transformer" in state_dict:
+                state_dict_ = state_dict["transformer"]
+            else:
+                # For backward compatibility.
+                state_dict_ = {}
+                for key in state_dict.keys():
+                    if "transformer." in key:
+                        state_dict_[key.split("transformer.")[1]] = state_dict[key]
+
+            # For backward compatibility.
+            state_dict_self_attention = {}
+            for key in state_dict_.keys():
+                if ".attention." in key:
+                    state_dict_self_attention[key.replace(".attention.", ".self_attention.")] = state_dict_[key]
+                else:
+                    state_dict_self_attention[key] = state_dict_[key]
+            state_dict_ = state_dict_self_attention
+
+            self.encoder.load_state_dict(state_dict_, strict=strict)
+
+        # Pooler.
+        if self.post_process:
+            if self.add_pooler:
+                assert "pooler" in state_dict, "could not find data for pooler in the checkpoint"
+                self.pooler.load_state_dict(state_dict[self._pooler_key], strict=strict)
+        # Decoder.
+        if self.add_decoder:
+            assert "decoder" in state_dict, "could not find data for pooler in the checkpoint"
+            self.decoder.load_state_dict(state_dict[self._decoder_key], strict=strict)
+
+
+def post_language_model_processing(lm_output, labels, logit_weights, parallel_output, fp16_lm_cross_entropy):
+    # Output.
+    output = parallel_lm_logits(lm_output, logit_weights, parallel_output)
+
+    if labels is None:
+        return output
+    else:
+        if fp16_lm_cross_entropy:
+            assert output.dtype == torch.half
+            loss = tensor_parallel.vocab_parallel_cross_entropy(output, labels)
+        else:
+            loss = tensor_parallel.vocab_parallel_cross_entropy(output.float(), labels)
+        return loss
+
+
+class GPTModel(MegatronModule):
+    """GPT-2 Language model."""
+
+    def __init__(self, num_tokentypes=0, parallel_output=True, pre_process=True, post_process=True):
+        super(GPTModel, self).__init__()
+        args = get_args()
+        self.parallel_output = parallel_output
+        self.pre_process = pre_process
+        self.post_process = post_process
+        self.fp16_lm_cross_entropy = args.fp16_lm_cross_entropy
+
+        self.language_model, self._language_model_key = get_language_model(
+            num_tokentypes=num_tokentypes,
+            add_pooler=False,
+            encoder_attn_mask_type=AttnMaskType.causal,
+            init_method=init_method_normal(args.init_method_std),
+            scaled_init_method=scaled_init_method_normal(args.init_method_std, args.num_layers),
+            pre_process=self.pre_process,
+            post_process=self.post_process,
+        )
+
+        self.initialize_word_embeddings(init_method_normal)
+
+    def set_input_tensor(self, input_tensor):
+        """See megatron.model.transformer.set_input_tensor()"""
+        self.language_model.set_input_tensor(input_tensor)
+
+    def forward(
+        self,
+        input_ids,
+        position_ids,
+        attention_mask,
+        labels=None,
+        tokentype_ids=None,
+        set_inference_key_value_memory=False,
+        inference_max_sequence_len=None,
+    ):
+
+        lm_output = self.language_model(
+            input_ids,
+            position_ids,
+            attention_mask,
+            set_inference_key_value_memory=set_inference_key_value_memory,
+            inference_max_sequence_len=inference_max_sequence_len,
+        )
+
+        if self.post_process:
+            return post_language_model_processing(
+                lm_output, labels, self.word_embeddings_weight(), self.parallel_output, self.fp16_lm_cross_entropy
+            )
+        else:
+            return lm_output
+
+    def state_dict_for_save_checkpoint(self, destination=None, prefix="", keep_vars=False):
+
+        state_dict_ = {}
+        state_dict_[self._language_model_key] = self.language_model.state_dict_for_save_checkpoint(
+            destination, prefix, keep_vars
+        )
+        # Save word_embeddings.
+        if self.post_process and not self.pre_process:
+            state_dict_[self._word_embeddings_for_head_key] = self.word_embeddings.state_dict(
+                destination, prefix, keep_vars
+            )
+        return state_dict_
+
+    def load_state_dict(self, state_dict, strict=True):
+        """Customized load."""
+
+        # Load word_embeddings.
+        if self.post_process and not self.pre_process:
+            self.word_embeddings.load_state_dict(state_dict[self._word_embeddings_for_head_key], strict=strict)
+        if self._language_model_key in state_dict:
+            state_dict = state_dict[self._language_model_key]
+        self.language_model.load_state_dict(state_dict, strict=strict)
+
+
+def gpt_model_provider(pre_process=True, post_process=True):
+    model = GPTModel(num_tokentypes=0, parallel_output=True, pre_process=pre_process, post_process=post_process)
+    return model
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/utils.py
new file mode 100644
index 0000000000..43d72c499e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/apex/transformer/utils.py
@@ -0,0 +1,47 @@
+"""Utility functions used by both `pipeline_parallel` and `tensor_parallel`"""
+import torch
+
+from apex.transformer import parallel_state
+
+
+def ensure_divisibility(numerator, denominator):
+    """Ensure that numerator is divisible by the denominator."""
+    assert numerator % denominator == 0, "{} is not divisible by {}".format(numerator, denominator)
+
+
+def divide(numerator, denominator):
+    """Ensure that numerator is divisible by the denominator and return
+    the division value."""
+    ensure_divisibility(numerator, denominator)
+    return numerator // denominator
+
+
+def split_tensor_into_1d_equal_chunks(tensor):
+    """Break a tensor into equal 1D chunks."""
+    data = tensor.view(-1)
+    partition_size = torch.numel(data) // parallel_state.get_tensor_model_parallel_world_size()
+    start_index = partition_size * parallel_state.get_tensor_model_parallel_rank()
+    end_index = start_index + partition_size
+    return data[start_index:end_index]
+
+
+def gather_split_1d_tensor(tensor):
+    """Opposite of above function, gather values from model parallel ranks."""
+    world_size = parallel_state.get_tensor_model_parallel_world_size()
+    numel = torch.numel(tensor)
+    numel_gathered = world_size * numel
+    gathered = torch.empty(numel_gathered, dtype=tensor.dtype, device=torch.cuda.current_device(), requires_grad=False)
+    chunks = [gathered[i * numel:(i + 1) * numel] for i in range(world_size)]
+    torch.distributed.all_gather(chunks, tensor, group=parallel_state.get_tensor_model_parallel_group())
+    return gathered
+
+
+# TODO(mkozuki): Rewrite this using `logging`.
+def rank_print(msg):
+    """Print the given msg with rank information"""
+    print(
+        f"tensor rank: {parallel_state.get_tensor_model_parallel_rank()}"
+        f"pipeline rank: {parallel_state.get_pipeline_model_parallel_rank()}, "
+        f"virtual pipeline rank: {parallel_state.get_virtual_pipeline_model_parallel_rank()}, "
+        f"data rank: {parallel_state.get_data_parallel_rank()} | {msg}"
+    )
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/amp_C_frontend.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/amp_C_frontend.cpp
new file mode 100644
index 0000000000..f1f74aa16f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/amp_C_frontend.cpp
@@ -0,0 +1,145 @@
+#include <torch/extension.h>
+
+void multi_tensor_scale_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  float scale);
+
+void multi_tensor_sgd_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  float wd,
+  float momentum,
+  float dampening,
+  float lr,
+  bool nesterov,
+  bool first_run,
+  bool wd_after_momentum,
+  float scale);
+
+void multi_tensor_axpby_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  float a,
+  float b,
+  int arg_to_check);
+
+std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::optional<bool> per_tensor_python);
+
+std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_scale_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  float scale,
+  at::optional<bool> per_tensor_python);
+
+void multi_tensor_lamb_stage1_cuda(
+    int chunk_size,
+    at::Tensor noop_flag,
+    std::vector<std::vector<at::Tensor>> tensor_lists,
+    at::Tensor per_tensor_decay,
+    const int step,
+    const float beta1,
+    const float beta2,
+    const float epsilon,
+    at::Tensor global_grad_norm,
+    const float max_global_grad_norm);
+
+void multi_tensor_lamb_stage2_cuda(
+    int chunk_size,
+    at::Tensor noop_flag,
+    std::vector<std::vector<at::Tensor>> tensor_lists,
+    at::Tensor per_tensor_param_norm,
+    at::Tensor per_tensor_update_norm,
+    const float lr,
+    const float weight_decay,
+    at::optional<bool> use_nvlamb_python);
+
+void multi_tensor_adam_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  const float lr,
+  const float beta1,
+  const float beta2,
+  const float epsilon,
+  const int step,
+  const int mode,
+  const int bias_correction,
+  const float weight_decay);
+
+
+void multi_tensor_adagrad_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  const float lr,
+  const float epsilon,
+  const int mode,
+  const float weight_decay);
+
+
+void multi_tensor_novograd_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor grad_norms,
+  const float lr,
+  const float beta1,
+  const float beta2,
+  const float epsilon,
+  const int step,
+  const int bias_correction,
+  const float weight_decay,
+  const int grad_averaging,
+  const int mode,
+  const int norm_type);
+
+void multi_tensor_lamb_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  const float lr,
+  const float beta1,
+  const float beta2,
+  const float epsilon,
+  const int step,
+  const int bias_correction,
+  const float weight_decay,
+  const int grad_averaging,
+  const int mode,
+  at::Tensor global_grad_norm,
+  const float max_grad_norm,
+  at::optional<bool> use_nvlamb_python);
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("multi_tensor_scale", &multi_tensor_scale_cuda,
+        "Fused overflow check + scale for a list of contiguous tensors");
+  m.def("multi_tensor_sgd", &multi_tensor_sgd_cuda,
+        "Fused SGD optimizer for list of contiguous tensors");
+  m.def("multi_tensor_axpby", &multi_tensor_axpby_cuda,
+        "out = a*x + b*y for a list of contiguous tensors");
+  m.def("multi_tensor_l2norm", &multi_tensor_l2norm_cuda,
+        "Computes L2 norm for a list of contiguous tensors");
+  m.def("multi_tensor_l2norm_scale", &multi_tensor_l2norm_scale_cuda,
+        "Computes L2 norm for a list of contiguous tensors and does scaling");
+  m.def("multi_tensor_lamb_stage1_cuda", &multi_tensor_lamb_stage1_cuda,
+        "Computes update part of LAMB optimizer");
+  m.def("multi_tensor_lamb_stage2_cuda", &multi_tensor_lamb_stage2_cuda,
+        "Completes application of gradient to parameters for LAMB optimizer");
+  m.def("multi_tensor_adam", &multi_tensor_adam_cuda,
+        "Compute and apply gradient update to parameters for Adam optimizer");
+  m.def("multi_tensor_adagrad", &multi_tensor_adagrad_cuda,
+        "Compute and apply gradient update to parameters for Adam optimizer");
+  m.def("multi_tensor_novograd", &multi_tensor_novograd_cuda,
+        "Compute and apply gradient update to parameters for Adam optimizer");
+  m.def("multi_tensor_lamb", &multi_tensor_lamb_cuda,
+        "Computes and apply update for LAMB optimizer");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/compat.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/compat.h
new file mode 100644
index 0000000000..acafb056b2
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/compat.h
@@ -0,0 +1,9 @@
+#ifndef TORCH_CHECK
+#define TORCH_CHECK AT_CHECK
+#endif
+
+#ifdef VERSION_GE_1_3
+#define DATA_PTR data_ptr
+#else
+#define DATA_PTR data
+#endif
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/flatten_unflatten.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/flatten_unflatten.cpp
new file mode 100644
index 0000000000..d49ce759f1
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/flatten_unflatten.cpp
@@ -0,0 +1,18 @@
+#include <torch/extension.h>
+#include <torch/csrc/utils/tensor_flatten.h>
+// https://github.com/pytorch/pytorch/blob/master/torch/csrc/utils/tensor_flatten.h
+
+at::Tensor flatten(std::vector<at::Tensor> tensors)
+{
+  return torch::utils::flatten_dense_tensors(tensors);
+}
+
+std::vector<at::Tensor> unflatten(at::Tensor flat, std::vector<at::Tensor> tensors)
+{
+  return torch::utils::unflatten_dense_tensors(flat, tensors);
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("flatten", &flatten, "Flatten dense tensors");
+  m.def("unflatten", &unflatten, "Unflatten dense tensors");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/fused_dense.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/fused_dense.cpp
new file mode 100644
index 0000000000..db5bd0d59d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/fused_dense.cpp
@@ -0,0 +1,192 @@
+#include <torch/extension.h>
+#include <torch/torch.h>
+#include <vector>
+
+#include <stdio.h>
+
+
+template <typename T>
+int linear_bias_forward_cuda(at::Tensor input, T *weight, at::Tensor bias, int in_features, int batch_size, int out_features, at::Tensor output, void *lt_workspace);
+
+template <typename T>
+int linear_bias_backward_cuda(T *input, T *weight, T *d_output, int in_features, int batch_size, int out_features, T *d_weight, T *d_bias, T *d_input,  void *lt_workspace);
+
+template <typename T>
+int linear_gelu_linear_forward_cuda(T *input, T *weight1, T *bias1, T *weight2, T *bias2, int in_features, int hidden_features, int batch_size, int out_features, T *output1, T *output2, T *gelu_in, void *lt_workspace) ;
+
+template <typename T>
+int linear_gelu_linear_backward_cuda(T *input, T *gelu_in, T *output1, T *weight1, T *weight2, T *d_output1, T *d_output2, int in_features, int batch_size, int hidden_features, int out_features, T *d_weight1, T *d_weight2, T *d_bias1, T *d_bias2, T *d_input, void *lt_workspace);
+
+at::Tensor linear_bias_forward(at::Tensor input, at::Tensor weight, at::Tensor bias) {
+
+  auto batch_size = input.size(0);
+  auto in_features = input.size(1);
+
+  int out_features = weight.size(0);
+
+  //auto reserved_size = get_mlp_reserved_space(batch_size, num_layers, output_features.data());
+
+  // create output/workspace tensor
+  auto out = at::empty({batch_size, out_features}, input.type());
+  //auto reserved_space = at::empty({reserved_size}, inputs[0].type());
+  // allocate fixed 4MB workspace for cublaslt for now, and this gets at least 4 MB
+  auto lt_workspace = at::empty({1 << 22}, input.type());
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "linear_bias_forward", [&] {
+    scalar_t* w_ptr = weight.data_ptr<scalar_t>();
+    scalar_t* b_ptr = bias.data_ptr<scalar_t>();
+    auto result = linear_bias_forward_cuda<scalar_t>(
+        input,
+        w_ptr,
+        bias,
+        in_features,
+        batch_size,
+        out_features,
+        out,
+        //out.data_ptr<scalar_t>(),
+       // reserved_space.data_ptr<scalar_t>(),
+        (void*) (lt_workspace.data_ptr<scalar_t>()));
+  });
+
+  return {out};
+}
+
+std::vector<at::Tensor> linear_bias_backward(at::Tensor input, at::Tensor weight, at::Tensor d_output) {
+
+  auto batch_size = input.size(0);
+  auto in_features = input.size(1);
+
+  int out_features = weight.size(0);
+
+  //auto reserved_size = get_mlp_reserved_space(batch_size, num_layers, output_features.data());
+
+  // create output/workspace tensor
+  auto d_weight = at::empty({out_features, in_features}, input.type());
+#if defined(CUBLAS_VERSION) && CUBLAS_VERSION < 11600
+  auto d_bias = d_output.view({-1, out_features}).sum(0, false);
+#else                                                                              
+  auto d_bias = at::empty({out_features}, input.type());
+#endif                                                                              
+  auto d_input = at::empty({batch_size, in_features}, input.type());
+  //auto reserved_space = at::empty({reserved_size}, inputs[0].type());
+  // allocate fixed 4MB workspace for cublaslt for now, and this gets at least 4 MB
+  auto lt_workspace = at::empty({1 << 22}, input.type());
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "linear_bias_backward", [&] {
+    scalar_t* w_ptr = weight.data_ptr<scalar_t>();
+    scalar_t* d_b_ptr = d_bias.data_ptr<scalar_t>();
+    auto result = linear_bias_backward_cuda<scalar_t>(
+        input.data_ptr<scalar_t>(),
+        w_ptr,
+        d_output.data_ptr<scalar_t>(),
+        in_features,
+        batch_size,
+        out_features,
+        d_weight.data_ptr<scalar_t>(),
+        d_bias.data_ptr<scalar_t>(),
+        d_input.data_ptr<scalar_t>(),
+       // reserved_space.data_ptr<scalar_t>(),
+        (void*) (lt_workspace.data_ptr<scalar_t>()));
+  });
+
+  return {d_input, d_weight, d_bias};
+}
+
+std::vector<at::Tensor> linear_gelu_linear_forward(at::Tensor input, at::Tensor weight1, at::Tensor bias1, at::Tensor weight2, at::Tensor bias2) {
+
+  auto batch_size = input.size(0);
+  auto in_features = input.size(1);
+
+  int hidden_features = weight1.size(0);
+  int out_features = weight2.size(0);
+
+  //auto reserved_size = get_mlp_reserved_space(batch_size, num_layers, output_features.data());
+
+  // create output/workspace tensor
+  auto output1 = at::empty({batch_size, hidden_features}, input.type());
+  auto gelu_in = at::empty({batch_size, hidden_features}, input.type());
+  auto output2 = at::empty({batch_size, out_features}, input.type());
+  //auto reserved_space = at::empty({reserved_size}, inputs[0].type());
+  // allocate fixed 4MB workspace for cublaslt for now, and this gets at least 4 MB
+  auto lt_workspace = at::empty({1 << 22}, input.type());
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "linear_gelu_linear_forward", [&] {
+    scalar_t* w1_ptr = weight1.data_ptr<scalar_t>();
+    scalar_t* b1_ptr = bias1.data_ptr<scalar_t>();
+    scalar_t* w2_ptr = weight2.data_ptr<scalar_t>();
+    scalar_t* b2_ptr = bias2.data_ptr<scalar_t>();
+    auto result = linear_gelu_linear_forward_cuda<scalar_t>(
+        input.data_ptr<scalar_t>(),
+        w1_ptr,
+        b1_ptr,
+	w2_ptr,
+	b2_ptr,
+        in_features,
+	hidden_features,
+        batch_size,
+        out_features,
+        output1.data_ptr<scalar_t>(),
+        output2.data_ptr<scalar_t>(),
+        gelu_in.data_ptr<scalar_t>(),
+       // reserved_space.data_ptr<scalar_t>(),
+        (void*) (lt_workspace.data_ptr<scalar_t>()));
+  });
+
+  return {output1, output2, gelu_in};
+}
+
+std::vector<at::Tensor> linear_gelu_linear_backward(at::Tensor input, at::Tensor gelu_in, at::Tensor output1, at::Tensor weight1, at::Tensor weight2, at::Tensor d_output2) {
+
+  auto batch_size = input.size(0);
+  auto in_features = input.size(1);
+
+  int hidden_features = weight1.size(0);
+  int out_features = weight2.size(0);
+
+  //auto reserved_size = get_mlp_reserved_space(batch_size, num_layers, output_features.data());
+
+  // create output/workspace tensor
+  auto d_weight1 = at::empty({hidden_features, in_features}, input.type());
+  auto d_weight2 = at::empty({out_features, hidden_features}, input.type());
+  auto d_bias1 = at::empty({hidden_features}, input.type());
+  auto d_bias2 = at::empty({out_features}, input.type());
+  auto d_input = at::empty({batch_size, in_features}, input.type());
+  auto d_output1 = at::empty({batch_size, hidden_features}, input.type());
+  //auto reserved_space = at::empty({reserved_size}, inputs[0].type());
+  // allocate fixed 4MB workspace for cublaslt for now, and this gets at least 4 MB
+  auto lt_workspace = at::empty({1 << 22}, input.type());
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(input.type(), "linear_bias_backward", [&] {
+    //scalar_t* w_ptr = weight.data_ptr<scalar_t>();
+    //scalar_t* d_b_ptr = d_bias.data_ptr<scalar_t>();
+    auto result = linear_gelu_linear_backward_cuda<scalar_t>(
+        input.data_ptr<scalar_t>(),
+        gelu_in.data_ptr<scalar_t>(),
+        output1.data_ptr<scalar_t>(),
+        weight1.data_ptr<scalar_t>(),
+        weight2.data_ptr<scalar_t>(),
+        d_output1.data_ptr<scalar_t>(),
+        d_output2.data_ptr<scalar_t>(),
+        in_features,
+        batch_size,
+        hidden_features,
+        out_features,
+        d_weight1.data_ptr<scalar_t>(),
+        d_weight2.data_ptr<scalar_t>(),
+        d_bias1.data_ptr<scalar_t>(),
+        d_bias2.data_ptr<scalar_t>(),
+        d_input.data_ptr<scalar_t>(),
+       // reserved_space.data_ptr<scalar_t>(),
+        (void*) (lt_workspace.data_ptr<scalar_t>()));
+  });
+
+  return {d_input, d_weight1, d_bias1, d_weight2, d_bias2};
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("linear_bias_forward", &linear_bias_forward, "linear bias forward");
+  m.def("linear_bias_backward", &linear_bias_backward, "linear bias backward");
+  m.def("linear_gelu_linear_forward", &linear_gelu_linear_forward, "linear gelu linear forward");
+  m.def("linear_gelu_linear_backward", &linear_gelu_linear_backward, "linear gelu linear backward");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/fused_dense_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/fused_dense_cuda.cu
new file mode 100644
index 0000000000..c12d264a18
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/fused_dense_cuda.cu
@@ -0,0 +1,1437 @@
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <torch/torch.h>
+
+/* Includes, cuda */
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000
+// includes cublaslt
+#include <cublasLt.h>
+#endif
+// FP64 Wrapper around cublas GEMMEx
+cublasStatus_t gemm_bias(
+    cublasHandle_t handle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float* alpha,
+    double* A,
+    int lda,
+    double* B,
+    int ldb,
+    const float* beta,
+    double* C,
+    int ldc) {
+  return cublasGemmEx(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      CUDA_R_64F,
+      lda,
+      B,
+      CUDA_R_64F,
+      ldb,
+      beta,
+      C,
+      CUDA_R_64F,
+      ldc,
+      CUDA_R_64F,
+      CUBLAS_GEMM_DEFAULT);
+}
+
+// FP32 Wrapper around cublas GEMMEx
+cublasStatus_t gemm_bias(
+    cublasHandle_t handle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float* alpha,
+    float* A,
+    int lda,
+    float* B,
+    int ldb,
+    const float* beta,
+    float* C,
+    int ldc) {
+  return cublasGemmEx(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      CUDA_R_32F,
+      lda,
+      B,
+      CUDA_R_32F,
+      ldb,
+      beta,
+      C,
+      CUDA_R_32F,
+      ldc,
+      CUDA_R_32F,
+      CUBLAS_GEMM_DEFAULT);
+}
+
+// FP16 Tensor core wrapper around cublas GEMMEx
+cublasStatus_t gemm_bias(
+    cublasHandle_t handle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float* alpha,
+    at::Half* A,
+    int lda,
+    at::Half* B,
+    int ldb,
+    const float* beta,
+    at::Half* C,
+    int ldc) {
+  return cublasGemmEx(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      CUDA_R_16F,
+      lda,
+      B,
+      CUDA_R_16F,
+      ldb,
+      beta,
+      C,
+      CUDA_R_16F,
+      ldc,
+      CUDA_R_32F,
+      CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+}
+
+
+#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11600
+
+
+int gemm_bias_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    at::Half* A,
+    int lda,
+    at::Half* B,
+    int ldb,
+    const float *beta, /* host pointer */
+    at::Half* C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    const void* bias) {
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  cublasLtMatmulDescOpaque_t operationDesc = {};
+  cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};
+  cublasLtMatmulPreferenceOpaque_t preference = {};
+
+  int returnedResults                             = 0;
+  cublasLtMatmulHeuristicResult_t heuristicResult = {};
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;
+
+  // Create operation descriptor; see cublasLtMatmulDescAttributes_t
+  // for details about defaults; here we just set the transforms for
+  // A and B.
+  status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (use_bias) {
+    status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias));
+    if (status != CUBLAS_STATUS_SUCCESS) {
+      goto CLEANUP;
+    }
+      epilogue = CUBLASLT_EPILOGUE_BIAS;
+  } 
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+
+  // Create matrix descriptors. Not setting any extra attributes.
+  status = cublasLtMatrixLayoutInit(
+    &Adesc, CUDA_R_16F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(
+    &Bdesc, CUDA_R_16F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_16F, m, n, ldc);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // Create preference handle; In general, extra attributes can be
+  // used here to disable tensor ops or to make sure algo selected
+  // will work with badly aligned A, B, C. However, for simplicity
+  // here we assume A,B,C are always well aligned (e.g., directly
+  // come from cudaMalloc)
+  status = cublasLtMatmulPreferenceInit(&preference);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulPreferenceSetAttribute(
+    &preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // We just need the best available heuristic to try and run matmul.
+  // There is no guarantee that this will work. For example, if A is
+  // badly aligned, you can request more (e.g. 32) algos and try to
+  // run them one by one until something works.
+  status = cublasLtMatmulAlgoGetHeuristic(
+    ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1, &heuristicResult, &returnedResults);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (returnedResults == 0) {
+    status = CUBLAS_STATUS_NOT_SUPPORTED;
+    goto CLEANUP;
+  }
+  status = cublasLtMatmul(ltHandle,
+                          &operationDesc,
+                          alpha,
+                          A,
+                          &Adesc,
+                          B,
+                          &Bdesc,
+                          beta,
+                          C,
+                          &Cdesc,
+                          C,
+                          &Cdesc,
+                          //&heuristicResult.algo,
+                          NULL,
+                          workspace,
+                          workspaceSize,
+                          stream);
+
+CLEANUP:
+  // Descriptors are no longer needed as all GPU work was already
+  // enqueued.
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;
+}
+
+
+
+
+
+
+
+int gemm_bias_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    double* A,
+    int lda,
+    double* B,
+    int ldb,
+    const float *beta, /* host pointer */
+    double* C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    const void* bias) {
+  return 1;
+}
+
+int gemm_bias_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    float *A,
+    int lda,
+    float *B,
+    int ldb,
+    const float *beta, /* host pointer */
+    float *C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    const void* bias) {
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  cublasLtMatmulDescOpaque_t operationDesc = {};
+  cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};
+  cublasLtMatmulPreferenceOpaque_t preference = {};
+
+  int returnedResults                             = 0;
+  cublasLtMatmulHeuristicResult_t heuristicResult = {};
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;
+
+  // Create operation descriptor; see cublasLtMatmulDescAttributes_t
+  // for details about defaults; here we just set the transforms for
+  // A and B.
+  status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (use_bias) {
+    status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias));
+    if (status != CUBLAS_STATUS_SUCCESS) {
+      goto CLEANUP;
+    }
+      epilogue = CUBLASLT_EPILOGUE_BIAS;
+  }
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+
+  // Create matrix descriptors. Not setting any extra attributes.
+  status = cublasLtMatrixLayoutInit(
+    &Adesc, CUDA_R_32F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(
+    &Bdesc, CUDA_R_32F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_32F, m, n, ldc);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // Create preference handle; In general, extra attributes can be
+  // used here to disable tensor ops or to make sure algo selected
+  // will work with badly aligned A, B, C. However, for simplicity
+  // here we assume A,B,C are always well aligned (e.g., directly
+  // come from cudaMalloc)
+  status = cublasLtMatmulPreferenceInit(&preference);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulPreferenceSetAttribute(
+    &preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // We just need the best available heuristic to try and run matmul.
+  // There is no guarantee that this will work. For example, if A is
+  // badly aligned, you can request more (e.g. 32) algos and try to
+  // run them one by one until something works.
+  status = cublasLtMatmulAlgoGetHeuristic(
+    ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1, &heuristicResult, &returnedResults);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (returnedResults == 0) {
+    status = CUBLAS_STATUS_NOT_SUPPORTED;
+    goto CLEANUP;
+  }
+
+  status = cublasLtMatmul(ltHandle,
+                          &operationDesc,
+                          alpha,
+                          A,
+                          &Adesc,
+                          B,
+                          &Bdesc,
+                          beta,
+                          C,
+                          &Cdesc,
+                          C,
+                          &Cdesc,
+                          &heuristicResult.algo,
+                          workspace,
+                          workspaceSize,
+                          stream);
+
+CLEANUP:
+  // Descriptors are no longer needed as all GPU work was already
+  // enqueued.
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;
+}
+
+
+
+int gemm_bias_gelu_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    at::Half* A,
+    int lda,
+    at::Half* B,
+    int ldb,
+    const float *beta, /* host pointer */
+    at::Half* C,
+    int64_t ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    const void* gelu_in,
+    const void* bias) {
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  cublasLtMatmulDescOpaque_t operationDesc = {};
+  cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};
+  cublasLtMatmulPreferenceOpaque_t preference = {};
+
+  int returnedResults                             = 0;
+  cublasLtMatmulHeuristicResult_t heuristicResult = {};
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_GELU_AUX;
+
+  // Create operation descriptor; see cublasLtMatmulDescAttributes_t
+  // for details about defaults; here we just set the transforms for
+  // A and B.
+  status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER, &gelu_in, sizeof(gelu_in));
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD, &ldc, sizeof(ldc));
+
+  if (use_bias) {
+    status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias));
+    if (status != CUBLAS_STATUS_SUCCESS) {
+      goto CLEANUP;
+    }
+      epilogue = CUBLASLT_EPILOGUE_GELU_AUX_BIAS;
+  } 
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+
+  // Create matrix descriptors. Not setting any extra attributes.
+  status = cublasLtMatrixLayoutInit(
+    &Adesc, CUDA_R_16F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(
+    &Bdesc, CUDA_R_16F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_16F, m, n, ldc);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // Create preference handle; In general, extra attributes can be
+  // used here to disable tensor ops or to make sure algo selected
+  // will work with badly aligned A, B, C. However, for simplicity
+  // here we assume A,B,C are always well aligned (e.g., directly
+  // come from cudaMalloc)
+  status = cublasLtMatmulPreferenceInit(&preference);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulPreferenceSetAttribute(
+    &preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // We just need the best available heuristic to try and run matmul.
+  // There is no guarantee that this will work. For example, if A is
+  // badly aligned, you can request more (e.g. 32) algos and try to
+  // run them one by one until something works.
+  status = cublasLtMatmulAlgoGetHeuristic(
+    ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1, &heuristicResult, &returnedResults);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (returnedResults == 0) {
+    status = CUBLAS_STATUS_NOT_SUPPORTED;
+    goto CLEANUP;
+  }
+  status = cublasLtMatmul(ltHandle,
+                          &operationDesc,
+                          alpha,
+                          A,
+                          &Adesc,
+                          B,
+                          &Bdesc,
+                          beta,
+                          C,
+                          &Cdesc,
+                          C,
+                          &Cdesc,
+                          //&heuristicResult.algo,
+                          NULL,
+                          workspace,
+                          workspaceSize,
+                          stream);
+
+CLEANUP:
+  // Descriptors are no longer needed as all GPU work was already
+  // enqueued.
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;
+}
+
+
+int gemm_bias_gelu_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    double* A,
+    int lda,
+    double* B,
+    int ldb,
+    const float *beta, /* host pointer */
+    double* C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    const void *gelu_in,
+    const void* bias) {
+  return 1;
+}
+
+
+int gemm_bias_gelu_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    float *A,
+    int lda,
+    float *B,
+    int ldb,
+    const float *beta, /* host pointer */
+    float *C,
+    int64_t ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    const void* gelu_in,
+    const void* bias) {
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  cublasLtMatmulDescOpaque_t operationDesc = {};
+  cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};
+  cublasLtMatmulPreferenceOpaque_t preference = {};
+
+  int returnedResults                             = 0;
+  cublasLtMatmulHeuristicResult_t heuristicResult = {};
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_GELU_AUX;
+
+  // Create operation descriptor; see cublasLtMatmulDescAttributes_t
+  // for details about defaults; here we just set the transforms for
+  // A and B.
+  status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER, &gelu_in, sizeof(gelu_in));
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD, &ldc, sizeof(ldc));
+
+  if (use_bias) {
+    status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias));
+    if (status != CUBLAS_STATUS_SUCCESS) {
+      goto CLEANUP;
+    }
+      epilogue = CUBLASLT_EPILOGUE_GELU_AUX_BIAS;
+  } 
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+
+  // Create matrix descriptors. Not setting any extra attributes.
+  status = cublasLtMatrixLayoutInit(
+    &Adesc, CUDA_R_32F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(
+    &Bdesc, CUDA_R_32F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_32F, m, n, ldc);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // Create preference handle; In general, extra attributes can be
+  // used here to disable tensor ops or to make sure algo selected
+  // will work with badly aligned A, B, C. However, for simplicity
+  // here we assume A,B,C are always well aligned (e.g., directly
+  // come from cudaMalloc)
+  status = cublasLtMatmulPreferenceInit(&preference);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulPreferenceSetAttribute(
+    &preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // We just need the best available heuristic to try and run matmul.
+  // There is no guarantee that this will work. For example, if A is
+  // badly aligned, you can request more (e.g. 32) algos and try to
+  // run them one by one until something works.
+  status = cublasLtMatmulAlgoGetHeuristic(
+    ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1, &heuristicResult, &returnedResults);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (returnedResults == 0) {
+    status = CUBLAS_STATUS_NOT_SUPPORTED;
+    goto CLEANUP;
+  }
+  status = cublasLtMatmul(ltHandle,
+                          &operationDesc,
+                          alpha,
+                          A,
+                          &Adesc,
+                          B,
+                          &Bdesc,
+                          beta,
+                          C,
+                          &Cdesc,
+                          C,
+                          &Cdesc,
+                          //&heuristicResult.algo,
+                          NULL,
+                          workspace,
+                          workspaceSize,
+                          stream);
+
+CLEANUP:
+  // Descriptors are no longer needed as all GPU work was already
+  // enqueued.
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;
+}
+
+
+
+int gemm_bgradb_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    at::Half* A,
+    int lda,
+    at::Half* B,
+    int ldb,
+    const float *beta, /* host pointer */
+    at::Half* C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    const void* bgrad) {
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  cublasLtMatmulDescOpaque_t operationDesc = {};
+  cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};
+  cublasLtMatmulPreferenceOpaque_t preference = {};
+
+  int returnedResults                             = 0;
+  cublasLtMatmulHeuristicResult_t heuristicResult = {};
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;
+
+  // Create operation descriptor; see cublasLtMatmulDescAttributes_t
+  // for details about defaults; here we just set the transforms for
+  // A and B.
+  status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (use_bias) {
+    status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bgrad, sizeof(bgrad));
+    if (status != CUBLAS_STATUS_SUCCESS) {
+      goto CLEANUP;
+    }
+      epilogue = CUBLASLT_EPILOGUE_BGRADB;
+  } 
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+
+  // Create matrix descriptors. Not setting any extra attributes.
+  status = cublasLtMatrixLayoutInit(
+    &Adesc, CUDA_R_16F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(
+    &Bdesc, CUDA_R_16F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_16F, m, n, ldc);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // Create preference handle; In general, extra attributes can be
+  // used here to disable tensor ops or to make sure algo selected
+  // will work with badly aligned A, B, C. However, for simplicity
+  // here we assume A,B,C are always well aligned (e.g., directly
+  // come from cudaMalloc)
+  status = cublasLtMatmulPreferenceInit(&preference);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulPreferenceSetAttribute(
+    &preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // We just need the best available heuristic to try and run matmul.
+  // There is no guarantee that this will work. For example, if A is
+  // badly aligned, you can request more (e.g. 32) algos and try to
+  // run them one by one until something works.
+  status = cublasLtMatmulAlgoGetHeuristic(
+    ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1, &heuristicResult, &returnedResults);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (returnedResults == 0) {
+    status = CUBLAS_STATUS_NOT_SUPPORTED;
+    goto CLEANUP;
+  }
+  status = cublasLtMatmul(ltHandle,
+                          &operationDesc,
+                          alpha,
+                          A,
+                          &Adesc,
+                          B,
+                          &Bdesc,
+                          beta,
+                          C,
+                          &Cdesc,
+                          C,
+                          &Cdesc,
+                          //&heuristicResult.algo,
+                          NULL,
+                          workspace,
+                          workspaceSize,
+                          stream);
+
+CLEANUP:
+  // Descriptors are no longer needed as all GPU work was already
+  // enqueued.
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;
+}
+
+
+
+
+
+
+
+int gemm_bgradb_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    double* A,
+    int lda,
+    double* B,
+    int ldb,
+    const float *beta, /* host pointer */
+    double* C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    const void* bgrad) {
+  return 1;
+}
+
+int gemm_bgradb_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    float *A,
+    int lda,
+    float *B,
+    int ldb,
+    const float *beta, /* host pointer */
+    float *C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    const void* bgrad) {
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  cublasLtMatmulDescOpaque_t operationDesc = {};
+  cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};
+  cublasLtMatmulPreferenceOpaque_t preference = {};
+
+  int returnedResults                             = 0;
+  cublasLtMatmulHeuristicResult_t heuristicResult = {};
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;
+
+  // Create operation descriptor; see cublasLtMatmulDescAttributes_t
+  // for details about defaults; here we just set the transforms for
+  // A and B.
+  status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (use_bias) {
+    status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bgrad, sizeof(bgrad));
+    if (status != CUBLAS_STATUS_SUCCESS) {
+      goto CLEANUP;
+    }
+      epilogue = CUBLASLT_EPILOGUE_BGRADB;
+  }
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+
+  // Create matrix descriptors. Not setting any extra attributes.
+  status = cublasLtMatrixLayoutInit(
+    &Adesc, CUDA_R_32F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(
+    &Bdesc, CUDA_R_32F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_32F, m, n, ldc);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // Create preference handle; In general, extra attributes can be
+  // used here to disable tensor ops or to make sure algo selected
+  // will work with badly aligned A, B, C. However, for simplicity
+  // here we assume A,B,C are always well aligned (e.g., directly
+  // come from cudaMalloc)
+  status = cublasLtMatmulPreferenceInit(&preference);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulPreferenceSetAttribute(
+    &preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // We just need the best available heuristic to try and run matmul.
+  // There is no guarantee that this will work. For example, if A is
+  // badly aligned, you can request more (e.g. 32) algos and try to
+  // run them one by one until something works.
+  status = cublasLtMatmulAlgoGetHeuristic(
+    ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1, &heuristicResult, &returnedResults);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (returnedResults == 0) {
+    status = CUBLAS_STATUS_NOT_SUPPORTED;
+    goto CLEANUP;
+  }
+
+  status = cublasLtMatmul(ltHandle,
+                          &operationDesc,
+                          alpha,
+                          A,
+                          &Adesc,
+                          B,
+                          &Bdesc,
+                          beta,
+                          C,
+                          &Cdesc,
+                          C,
+                          &Cdesc,
+                          &heuristicResult.algo,
+                          workspace,
+                          workspaceSize,
+                          stream);
+
+CLEANUP:
+  // Descriptors are no longer needed as all GPU work was already
+  // enqueued.
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;
+}
+
+
+int gemm_dgelu_bgradb_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    at::Half* A,
+    int lda,
+    at::Half* B,
+    int ldb,
+    const float *beta, /* host pointer */
+    at::Half* C,
+    int64_t ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    const void *gelu_in,
+    const void *bgrad) {
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  cublasLtMatmulDescOpaque_t operationDesc = {};
+  cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};
+  cublasLtMatmulPreferenceOpaque_t preference = {};
+
+  int returnedResults                             = 0;
+  cublasLtMatmulHeuristicResult_t heuristicResult = {};
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DGELU_BGRAD;
+
+  // Create operation descriptor; see cublasLtMatmulDescAttributes_t
+  // for details about defaults; here we just set the transforms for
+  // A and B.
+  status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bgrad, sizeof(bgrad));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER, &gelu_in, sizeof(gelu_in));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD, &ldc, sizeof(ldc));
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+
+  // Create matrix descriptors. Not setting any extra attributes.
+  status = cublasLtMatrixLayoutInit(
+    &Adesc, CUDA_R_16F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(
+    &Bdesc, CUDA_R_16F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_16F, m, n, ldc);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // Create preference handle; In general, extra attributes can be
+  // used here to disable tensor ops or to make sure algo selected
+  // will work with badly aligned A, B, C. However, for simplicity
+  // here we assume A,B,C are always well aligned (e.g., directly
+  // come from cudaMalloc)
+  status = cublasLtMatmulPreferenceInit(&preference);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulPreferenceSetAttribute(
+    &preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // We just need the best available heuristic to try and run matmul.
+  // There is no guarantee that this will work. For example, if A is
+  // badly aligned, you can request more (e.g. 32) algos and try to
+  // run them one by one until something works.
+  status = cublasLtMatmulAlgoGetHeuristic(
+    ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1, &heuristicResult, &returnedResults);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (returnedResults == 0) {
+    status = CUBLAS_STATUS_NOT_SUPPORTED;
+    goto CLEANUP;
+  }
+  status = cublasLtMatmul(ltHandle,
+                          &operationDesc,
+                          alpha,
+                          A,
+                          &Adesc,
+                          B,
+                          &Bdesc,
+                          beta,
+                          C,
+                          &Cdesc,
+                          C,
+                          &Cdesc,
+                          //&heuristicResult.algo,
+                          NULL,
+                          workspace,
+                          workspaceSize,
+                          stream);
+
+CLEANUP:
+  // Descriptors are no longer needed as all GPU work was already
+  // enqueued.
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;
+}
+
+int gemm_dgelu_bgradb_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    double *A,
+    int lda,
+    double *B,
+    int ldb,
+    const float *beta, /* host pointer */
+    double *C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    const void *gelu_in,
+    const void *bgrad) {
+    return 1;
+}
+
+int gemm_dgelu_bgradb_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    const float *alpha, /* host pointer */
+    float *A,
+    int lda,
+    float *B,
+    int ldb,
+    const float *beta, /* host pointer */
+    float *C,
+    int64_t ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    const void *gelu_in,
+    const void *bgrad) {
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  cublasLtMatmulDescOpaque_t operationDesc = {};
+  cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};
+  cublasLtMatmulPreferenceOpaque_t preference = {};
+
+  int returnedResults                             = 0;
+  cublasLtMatmulHeuristicResult_t heuristicResult = {};
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DGELU_BGRAD;
+
+  // Create operation descriptor; see cublasLtMatmulDescAttributes_t
+  // for details about defaults; here we just set the transforms for
+  // A and B.
+  status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bgrad, sizeof(bgrad));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_POINTER, &gelu_in, sizeof(gelu_in));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE_AUX_LD, &ldc, sizeof(ldc));
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+
+  // Create matrix descriptors. Not setting any extra attributes.
+  status = cublasLtMatrixLayoutInit(
+    &Adesc, CUDA_R_32F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(
+    &Bdesc, CUDA_R_32F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_32F, m, n, ldc);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // Create preference handle; In general, extra attributes can be
+  // used here to disable tensor ops or to make sure algo selected
+  // will work with badly aligned A, B, C. However, for simplicity
+  // here we assume A,B,C are always well aligned (e.g., directly
+  // come from cudaMalloc)
+  status = cublasLtMatmulPreferenceInit(&preference);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulPreferenceSetAttribute(
+    &preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // We just need the best available heuristic to try and run matmul.
+  // There is no guarantee that this will work. For example, if A is
+  // badly aligned, you can request more (e.g. 32) algos and try to
+  // run them one by one until something works.
+  status = cublasLtMatmulAlgoGetHeuristic(
+    ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1, &heuristicResult, &returnedResults);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (returnedResults == 0) {
+    status = CUBLAS_STATUS_NOT_SUPPORTED;
+    goto CLEANUP;
+  }
+  status = cublasLtMatmul(ltHandle,
+                          &operationDesc,
+                          alpha,
+                          A,
+                          &Adesc,
+                          B,
+                          &Bdesc,
+                          beta,
+                          C,
+                          &Cdesc,
+                          C,
+                          &Cdesc,
+                          //&heuristicResult.algo,
+                          NULL,
+                          workspace,
+                          workspaceSize,
+                          stream);
+
+CLEANUP:
+  // Descriptors are no longer needed as all GPU work was already
+  // enqueued.
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;
+}
+
+#endif
+
+template <typename T>
+int linear_bias_forward_cuda(at::Tensor input, T *weight, at::Tensor bias, int in_features, int batch_size, int out_features, at::Tensor output, void *lt_workspace) {
+    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+    // Get the stream from cublas handle to reuse for biasReLU kernel.
+    cudaStream_t stream;
+    cublasGetStream(handle, &stream);
+    const float alpha          = 1.0;
+    const float beta_zero       = 0.0;
+    const float beta_one       = 1.0;
+    int status = 1;
+#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11600
+    status = gemm_bias_lt(
+    (cublasLtHandle_t)handle,
+    CUBLAS_OP_T,
+    CUBLAS_OP_N,
+    out_features,
+    batch_size,
+    in_features,
+    &alpha, /* host pointer */
+    weight,
+    in_features,
+    input.data_ptr<T>(),
+    in_features,
+    &beta_zero, /* host pointer */
+    output.data_ptr<T>(),
+    out_features,
+    lt_workspace,
+    1 << 22,
+    stream,
+    true,
+    static_cast<const void*>(bias.data_ptr<T>()));
+#endif
+    if (status != 0){
+        output.copy_(bias);
+        status = gemm_bias(
+          handle,
+          CUBLAS_OP_T,
+          CUBLAS_OP_N,
+          out_features,
+          batch_size,
+          in_features,
+          &alpha,
+          weight,
+          in_features,
+          input.data_ptr<T>(),
+          in_features,
+          &beta_one,
+          output.data_ptr<T>(),
+          out_features);
+    }
+    return status;
+}
+
+    
+template <typename T>
+int linear_bias_backward_cuda(T *input, T *weight, T *d_output, int in_features, int batch_size, int out_features, T *d_weight, T *d_bias, T *d_input,  void *lt_workspace) {
+    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+    // Get the stream from cublas handle to reuse for biasReLU kernel.
+    cudaStream_t stream;
+    cublasGetStream(handle, &stream);
+    const float alpha          = 1.0;
+    const float beta_zero       = 0.0;
+    int status = 1;
+#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11600
+    status = gemm_bgradb_lt(
+    (cublasLtHandle_t)handle,
+    CUBLAS_OP_N,
+    CUBLAS_OP_T,
+    in_features,
+    out_features,
+    batch_size,
+    &alpha, /* host pointer */
+    input,
+    in_features,
+    d_output,
+    out_features,
+    &beta_zero, /* host pointer */
+    d_weight,
+    in_features,
+    lt_workspace,
+    1 << 22,
+    stream,
+    true,
+    static_cast<const void*>(d_bias));
+#endif
+    
+
+    if (status != 0){
+    
+        status = gemm_bias(
+          handle,
+          CUBLAS_OP_N,
+          CUBLAS_OP_T,
+          in_features,
+          out_features,
+          batch_size,
+          &alpha,
+          input,
+          in_features,
+          d_output,
+          out_features,
+          &beta_zero,
+          d_weight,
+          in_features);
+    }
+    
+    status = gemm_bias(
+      handle,
+      CUBLAS_OP_N,
+      CUBLAS_OP_N,
+      in_features,
+      batch_size,
+      out_features,
+      &alpha,
+      weight,
+      in_features,
+      d_output,
+      out_features,
+      &beta_zero,
+      d_input,
+      in_features);
+    return status;
+
+}
+
+template <typename T>
+int linear_gelu_linear_forward_cuda(T *input, T *weight1, T *bias1, T *weight2, T *bias2, int in_features, int hidden_features, int batch_size, int out_features, T *output1, T *output2, T *gelu_in, void *lt_workspace) {
+    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+    // Get the stream from cublas handle to reuse for biasReLU kernel.
+    cudaStream_t stream;
+    cublasGetStream(handle, &stream);
+    const float alpha          = 1.0;
+    const float beta_zero       = 0.0;
+    int status = 1;
+#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11600
+    status = gemm_bias_gelu_lt(
+    (cublasLtHandle_t)handle,
+    CUBLAS_OP_T,
+    CUBLAS_OP_N,
+    hidden_features,
+    batch_size,
+    in_features,
+    &alpha, /* host pointer */
+    weight1,
+    in_features,
+    input,
+    in_features,
+    &beta_zero, /* host pointer */
+    output1,
+    hidden_features,
+    lt_workspace,
+    1 << 22,
+    stream,
+    true,
+    static_cast<const void*>(gelu_in),
+    static_cast<const void*>(bias1));
+    status = gemm_bias_lt(
+    (cublasLtHandle_t)handle,
+    CUBLAS_OP_T,
+    CUBLAS_OP_N,
+    out_features,
+    batch_size,
+    hidden_features,
+    &alpha, /* host pointer */
+    weight2,
+    hidden_features,
+    output1,
+    hidden_features,
+    &beta_zero, /* host pointer */
+    output2,
+    out_features,
+    lt_workspace,
+    1 << 22,
+    stream,
+    true,
+    static_cast<const void*>(bias2));
+    return status;
+#else 
+    return 1;
+#endif
+}
+
+template <typename T>
+int linear_gelu_linear_backward_cuda(T *input, T *gelu_in, T *output1, T *weight1, T *weight2, T *d_output1, T *d_output2, int in_features, int batch_size, int hidden_features, int out_features, T *d_weight1, T *d_weight2, T *d_bias1, T *d_bias2, T *d_input, void *lt_workspace) {
+    cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+    // Get the stream from cublas handle to reuse for biasReLU kernel.
+    cudaStream_t stream;
+    cublasGetStream(handle, &stream);
+    const float alpha          = 1.0;
+    const float beta_zero       = 0.0;
+    int status = 1;
+#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11600
+//wgrad for first gemm
+    status = gemm_bgradb_lt(
+    (cublasLtHandle_t)handle,
+    CUBLAS_OP_N,
+    CUBLAS_OP_T,
+    hidden_features,
+    out_features,
+    batch_size,
+    &alpha, /* host pointer */
+    output1,
+    hidden_features,
+    d_output2,
+    out_features,
+    &beta_zero, /* host pointer */
+    d_weight2,
+    hidden_features,
+    lt_workspace,
+    1 << 22,
+    stream,
+    true,
+    static_cast<const void*>(d_bias2));
+//dgrad for second GEMM
+    status = gemm_dgelu_bgradb_lt(
+    (cublasLtHandle_t)handle,
+    CUBLAS_OP_N,
+    CUBLAS_OP_N,
+    hidden_features,
+    batch_size,
+    out_features,
+    &alpha, /* host pointer */
+    weight2,
+    hidden_features,
+    d_output2,
+    out_features,
+    &beta_zero, /* host pointer */
+    d_output1,
+    hidden_features,
+    lt_workspace,
+    1 << 22,
+    stream,
+    static_cast<const void*>(gelu_in),
+    static_cast<const void*>(d_bias1));
+//wgrad for the first GEMM
+    status = gemm_bias(
+      handle,
+      CUBLAS_OP_N,
+      CUBLAS_OP_T,
+      in_features,
+      hidden_features,
+      batch_size,
+      &alpha,
+      input,
+      in_features,
+      d_output1,
+      hidden_features,
+      &beta_zero,
+      d_weight1,
+      in_features);
+
+//dgrad for the first GEMM
+    status = gemm_bias(
+      handle,
+      CUBLAS_OP_N,
+      CUBLAS_OP_N,
+      in_features,
+      batch_size,
+      hidden_features,
+      &alpha,
+      weight1,
+      in_features,
+      d_output1,
+      hidden_features,
+      &beta_zero,
+      d_input,
+      in_features);
+#endif
+    return status;
+
+}
+
+
+template int linear_bias_forward_cuda<at::Half>(at::Tensor input, at::Half *weight, at::Tensor bias, int in_features, int batch_size, int out_features, at::Tensor output, void *lt_workspace);
+
+template int linear_bias_forward_cuda<float>(at::Tensor input, float *weight, at::Tensor bias, int in_features, int batch_size, int out_features, at::Tensor output, void *lt_workspace);
+
+template int linear_bias_forward_cuda<double>(at::Tensor input, double *weight, at::Tensor bias, int in_features, int batch_size, int out_features, at::Tensor output, void *lt_workspace);
+
+template int linear_bias_backward_cuda<at::Half>(at::Half *input, at::Half *weight, at::Half *d_output, int in_features, int batch_size, int out_features, at::Half *d_weight, at::Half *d_bias, at::Half *d_input,  void *lt_workspace) ;
+
+template int linear_bias_backward_cuda<float>(float *input, float *weight, float *d_output, int in_features, int batch_size, int out_features, float *d_weight, float *d_bias, float *d_input,  void *lt_workspace) ;
+
+template int linear_bias_backward_cuda<double>(double *input, double *weight, double *d_output, int in_features, int batch_size, int out_features, double *d_weight, double *d_bias, double *d_input,  void *lt_workspace) ;
+
+
+template int linear_gelu_linear_forward_cuda<at::Half>(at::Half *input, at::Half *weight1, at::Half *bias1, at::Half *weight2, at::Half *bias2, int in_features, int hidden_features, int batch_size, int out_features, at::Half *output1, at::Half *output2, at::Half *gelu_in, void *lt_workspace) ;
+
+template int linear_gelu_linear_forward_cuda<float>(float *input, float *weight1, float *bias1, float *weight2, float *bias2, int in_features, int hidden_features, int batch_size, int out_features, float *output1, float *output2, float *gelu_in, void *lt_workspace);
+
+template int linear_gelu_linear_forward_cuda<double>(double *input, double *weight1, double *bias1, double *weight2, double *bias2, int in_features, int hidden_features, int batch_size, int out_features, double *output1, double *output2, double *gelu_in, void *lt_workspace) ;
+
+template int linear_gelu_linear_backward_cuda<at::Half>(at::Half *input, at::Half *gelu_in, at::Half *output1, at::Half *weight1, at::Half *weight2, at::Half *d_output1, at::Half *d_output2, int in_features, int batch_size, int hidden_features, int out_features, at::Half *d_weight1, at::Half *d_weight2, at::Half *d_bias1, at::Half *d_bias2, at::Half *d_input, void *lt_workspace);
+
+template int linear_gelu_linear_backward_cuda<float>(float *input, float *gelu_in, float *output1, float *weight1, float *weight2, float *d_output1, float *d_output2, int in_features, int batch_size, int hidden_features, int out_features, float *d_weight1, float *d_weight2, float *d_bias1, float *d_bias2, float *d_input, void *lt_workspace);
+
+template int linear_gelu_linear_backward_cuda<double>(double *input, double *gelu_in, double *output1, double *weight1, double *weight2, double *d_output1, double *d_output2, int in_features, int batch_size, int hidden_features, int out_features, double *d_weight1, double *d_weight2, double *d_bias1, double *d_bias2, double *d_input, void *lt_workspace);
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/layer_norm_cuda.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/layer_norm_cuda.cpp
new file mode 100644
index 0000000000..df5d4b4040
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/layer_norm_cuda.cpp
@@ -0,0 +1,267 @@
+#include <torch/extension.h>
+#include <vector>
+#include <cassert>
+#include "compat.h"
+
+namespace {
+void compute_n1_n2(
+    at::Tensor input,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    int& n1,
+    int& n2)
+{
+    int idiff = input.ndimension() - normalized_shape.size();
+    n2 = 1;
+    for (int i = 0;  i < (int)normalized_shape.size();  ++i) {
+	    assert( input.sizes()[i+idiff] == normalized_shape[i] );
+	    n2 *= normalized_shape[i];
+    }
+    n1 = 1;
+    for (int i = 0;  i < idiff;  ++i) {
+	    n1 *= input.sizes()[i];
+    }
+}
+
+void check_args(
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    at::Tensor gamma,
+    at::Tensor beta
+    )
+{
+    TORCH_CHECK(!gamma.defined() || gamma.sizes().equals(normalized_shape));
+    TORCH_CHECK(!beta.defined() || beta.sizes().equals(normalized_shape));
+}
+
+void check_args(
+    at::Tensor input,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    int& n1,
+    int& n2
+    )
+{
+    int64_t normalized_ndim = normalized_shape.size();
+
+    if (normalized_ndim < 1) {
+      std::stringstream ss;
+      ss << "Expected normalized_shape to be at least 1-dimensional, i.e., "
+         << "containing at least one element, but got normalized_shape="
+         << normalized_shape;
+      throw std::runtime_error(ss.str());
+    }
+
+    auto input_shape = input.sizes();
+    auto input_ndim = input.dim();
+
+    if (input_ndim < normalized_ndim ||
+        !input_shape.slice(input_ndim - normalized_ndim).equals(normalized_shape)) {
+      std::stringstream ss;
+      ss << "Given normalized_shape=" << normalized_shape
+         << ", expected input with shape [*";
+      for (auto size : normalized_shape) {
+        ss << ", " << size;
+      }
+      ss << "], but got input of size" << input_shape;
+      throw std::runtime_error(ss.str());
+    }
+
+    compute_n1_n2(input,normalized_shape,n1,n2);
+}
+
+
+void check_args(
+    at::Tensor input,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    at::Tensor gamma,
+    at::Tensor beta,
+    int& n1,
+    int& n2
+    )
+{
+    check_args(input,normalized_shape,n1,n2);
+    check_args(normalized_shape,gamma,beta);
+}
+}
+
+void cuda_layer_norm(
+    at::Tensor* output,
+    at::Tensor* mean,
+    at::Tensor* invvar,
+    at::Tensor* input,
+    int n1,
+    int n2,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    at::Tensor* gamma,
+    at::Tensor* beta,
+    double epsilon);
+
+#define CHECK_CUDA(x) TORCH_CHECK(x.type().is_cuda(), #x " must be a CUDA tensor")
+#define CHECK_CONTIGUOUS(x) TORCH_CHECK(x.is_contiguous(), #x " must be contiguous")
+#define CHECK_INPUT(x) CHECK_CUDA(x); CHECK_CONTIGUOUS(x)
+
+std::vector<at::Tensor> layer_norm(
+    at::Tensor input,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    double epsilon) {
+  CHECK_INPUT(input);
+  int n1,n2;
+  check_args(input,normalized_shape,n1,n2);
+  at::Tensor output = at::empty_like(input);
+  at::Tensor mean = at::empty({n1}, input.options().dtype(input.scalar_type()==at::ScalarType::Half || input.scalar_type()==at::ScalarType::BFloat16 ? at::ScalarType::Float : input.scalar_type()));
+  at::Tensor invvar = at::empty_like(mean);
+  cuda_layer_norm(&output,&mean,&invvar,&input,n1,n2,
+      normalized_shape,NULL,NULL,epsilon);
+  return {output, mean, invvar};
+}
+
+std::vector<at::Tensor> layer_norm_affine(
+    at::Tensor input,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    at::Tensor gamma,
+    at::Tensor beta,
+    double epsilon) {
+  CHECK_INPUT(input);
+  CHECK_INPUT(gamma);
+  CHECK_INPUT(beta);
+  int n1,n2;
+  check_args(input,normalized_shape,gamma,beta,n1,n2);
+  at::Tensor output = at::empty_like(input);
+  const auto stats_dtype = (input.scalar_type() == at::ScalarType::Half || input.scalar_type() == at::ScalarType::BFloat16) ? at::ScalarType::Float : input.scalar_type();
+  at::Tensor mean = at::empty({n1}, input.options().dtype(stats_dtype));
+  at::Tensor invvar = at::empty_like(mean);
+  cuda_layer_norm(&output,&mean,&invvar,&input,n1,n2,
+      normalized_shape,&gamma,&beta,epsilon);
+  return {output, mean, invvar};
+}
+
+std::vector<at::Tensor> layer_norm_affine_mixed_dtypes(
+    at::Tensor input,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    at::Tensor gamma,
+    at::Tensor beta,
+    double epsilon) {
+  CHECK_INPUT(input);
+  int n1, n2;
+  check_args(input, normalized_shape, n1, n2);
+  at::Tensor output = at::empty_like(input, gamma.options().dtype(gamma.scalar_type()));
+  at::Tensor mean = at::empty({n1}, input.options().dtype(input.scalar_type() == at::ScalarType::Half || input.scalar_type() == at::ScalarType::BFloat16 ? at::ScalarType::Float : input.scalar_type()));
+  at::Tensor invvar = at::empty_like(mean);
+   cuda_layer_norm(&output, &mean, &invvar, &input, n1, n2,
+      normalized_shape, &gamma, &beta, epsilon);
+  return {output, mean, invvar};
+}
+
+void cuda_layer_norm_gradient(
+    at::Tensor* dout,
+    at::Tensor* mean,
+    at::Tensor* invvar,
+    at::Tensor* input,
+    int n1,
+    int n2,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    at::Tensor* gamma,
+    at::Tensor* beta,
+    double epsilon,
+    at::Tensor* grad_input,
+    at::Tensor* grad_gamma,
+    at::Tensor* grad_beta
+    );
+
+at::Tensor layer_norm_gradient(
+    at::Tensor dout,
+    at::Tensor mean,
+    at::Tensor invvar,
+    at::Tensor input,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    double epsilon) {
+  CHECK_INPUT(dout);
+  CHECK_INPUT(mean);
+  CHECK_INPUT(invvar);
+  CHECK_INPUT(input);
+  int n1,n2;
+  check_args(input,normalized_shape,n1,n2);
+  at::Tensor grad_input = at::empty_like(input);
+  cuda_layer_norm_gradient(&dout,&mean,&invvar,&input,n1,n2,
+      normalized_shape,NULL,NULL,epsilon,
+      &grad_input,NULL,NULL);
+  return grad_input;
+}
+
+std::vector<at::Tensor> layer_norm_gradient_affine(
+    at::Tensor dout,
+    at::Tensor mean,
+    at::Tensor invvar,
+    at::Tensor input,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    at::Tensor gamma,
+    at::Tensor beta,
+    double epsilon) {
+  CHECK_INPUT(dout);
+  CHECK_INPUT(mean);
+  CHECK_INPUT(invvar);
+  CHECK_INPUT(input);
+  CHECK_INPUT(gamma);
+  CHECK_INPUT(beta);
+  int n1,n2;
+  check_args(input,normalized_shape,gamma,beta,n1,n2);
+  at::Tensor grad_input = at::empty_like(input);
+  at::Tensor grad_gamma = at::empty_like(gamma);
+  at::Tensor grad_beta = at::empty_like(beta);
+  cuda_layer_norm_gradient(&dout,&mean,&invvar,&input,n1,n2,
+      normalized_shape,&gamma,&beta,epsilon,
+      &grad_input,&grad_gamma,&grad_beta);
+  return {grad_input, grad_gamma, grad_beta};
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward_affine", &layer_norm_affine, "LayerNorm forward (CUDA)");
+  m.def("forward", &layer_norm, "LayerNorm forward (CUDA)");
+  m.def("backward_affine", &layer_norm_gradient_affine, "LayerNorm backward (CUDA)");
+  m.def("backward", &layer_norm_gradient, "LayerNorm backward (CUDA)");
+
+  m.def("forward_affine_mixed_dtypes", &layer_norm_affine_mixed_dtypes, "LayerNorm forward with mixed dtypes (CUDA) compatible with Megatron's implementation");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/layer_norm_cuda_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/layer_norm_cuda_kernel.cu
new file mode 100644
index 0000000000..29c7019a8b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/layer_norm_cuda_kernel.cu
@@ -0,0 +1,830 @@
+#include "ATen/ATen.h"
+#include "ATen/AccumulateType.h"
+#include "ATen/cuda/CUDAContext.h"
+#include "ATen/cuda/DeviceUtils.cuh"
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include "type_shim.h"
+
+template<typename U> __device__
+void cuWelfordOnlineSum(
+  const U curr,
+  U& mu,
+  U& sigma2,
+  U& count)
+{
+  count = count + U(1);
+  U delta = curr - mu;
+  U lmean = mu + delta / count;
+  mu = lmean;
+  U delta2 = curr - lmean;
+  sigma2 = sigma2 + delta * delta2;
+}
+
+template<typename U> __device__
+void cuChanOnlineSum(
+  const U muB,
+  const U sigma2B,
+  const U countB,
+  U& mu,
+  U& sigma2,
+  U& count)
+{
+  U delta = muB - mu;
+  U nA = count;
+  U nB = countB;
+  count = count + countB;
+  U nX = count;
+  if (nX > U(0)) {
+    nA = nA / nX;
+    nB = nB / nX;
+    mu = nA*mu + nB*muB;
+    sigma2 = sigma2 + sigma2B + delta * delta * nA * nB * nX;
+  } else {
+    mu = U(0);
+    sigma2 = U(0);
+  }
+}
+
+template<typename T, typename U> __device__
+void cuWelfordMuSigma2(
+  const T* __restrict__ vals,
+  const int n1,
+  const int n2,
+  const int i1,
+  U& mu,
+  U& sigma2,
+  U* buf)
+{
+  // Assumptions:
+  // 1) blockDim.x == warpSize
+  // 2) Tensor is contiguous
+  // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.
+  //
+  // compute variance and mean over n2
+  U count = U(0);
+  mu= U(0);
+  sigma2 = U(0);
+  if (i1 < n1) {
+    // one warp normalizes one n1 index,
+    // synchronization is implicit
+    // initialize with standard Welford algorithm
+    const int numx = blockDim.x * blockDim.y;
+    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
+    const T* lvals = vals + i1*n2;
+    int l = 4*thrx;
+    for (;  l+3 < n2;  l+=4*numx) {
+      for (int k = 0;  k < 4;  ++k) {
+        U curr = static_cast<U>(lvals[l+k]);
+        cuWelfordOnlineSum<U>(curr,mu,sigma2,count);
+      }
+    }
+    for (;  l < n2;  ++l) {
+      U curr = static_cast<U>(lvals[l]);
+      cuWelfordOnlineSum<U>(curr,mu,sigma2,count);
+    }
+    // intra-warp reductions
+    for (int l = 0;  l <= 4;  ++l) {
+      int srcLaneB = (threadIdx.x+(1<<l))&31;
+      U muB = WARP_SHFL(mu, srcLaneB);
+      U countB = WARP_SHFL(count, srcLaneB);
+      U sigma2B = WARP_SHFL(sigma2, srcLaneB);
+      cuChanOnlineSum<U>(muB,sigma2B,countB,mu,sigma2,count);
+    }
+    // threadIdx.x == 0 has correct values for each warp
+    // inter-warp reductions
+    if (blockDim.y > 1) {
+      U* ubuf = (U*)buf;
+      U* ibuf = (U*)(ubuf + blockDim.y);
+      for (int offset = blockDim.y/2;  offset > 0;  offset /= 2) {
+        // upper half of warps write to shared
+        if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2*offset) {
+          const int wrt_y = threadIdx.y - offset;
+          ubuf[2*wrt_y] = mu;
+          ubuf[2*wrt_y+1] = sigma2;
+          ibuf[wrt_y] = count;
+        }
+        __syncthreads();
+        // lower half merges
+        if (threadIdx.x == 0 && threadIdx.y < offset) {
+          U muB = ubuf[2*threadIdx.y];
+          U sigma2B = ubuf[2*threadIdx.y+1];
+          U countB = ibuf[threadIdx.y];
+          cuChanOnlineSum<U>(muB,sigma2B,countB,mu,sigma2,count);
+        }
+        __syncthreads();
+      }
+      // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values
+      if (threadIdx.x == 0 && threadIdx.y == 0) {
+        ubuf[0] = mu;
+        ubuf[1] = sigma2;
+      }
+      __syncthreads();
+      mu = ubuf[0];
+      sigma2 = ubuf[1]/U(n2);
+      // don't care about final value of count, we know count == n2
+    } else {
+      mu = WARP_SHFL(mu, 0);
+      sigma2 = WARP_SHFL(sigma2/U(n2), 0);
+    }
+  }
+}
+
+template<> __device__
+void cuWelfordMuSigma2(
+  const at::Half* __restrict__ vals,
+  const int n1,
+  const int n2,
+  const int i1,
+  float& mu,
+  float& sigma2,
+  float* buf)
+{
+  // Assumptions:
+  // 1) blockDim.x == warpSize
+  // 2) Tensor is contiguous
+  // 3) 2*blockDim.y*sizeof(U)+blockDim.y*sizeof(int) shared memory available.
+  //
+  // compute variance and mean over n2
+  float count = 0.0f;
+  mu= float(0);
+  sigma2 = float(0);
+  if (i1 < n1) {
+    // one warp normalizes one n1 index,
+    // synchronization is implicit
+    // initialize with standard Welford algorithm
+    const int numx = blockDim.x * blockDim.y;
+    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
+    const at::Half* lvals = vals + i1*n2;
+    int l = 8*thrx;
+    if ((((size_t)lvals)&3) != 0) {
+      // 16 bit alignment
+      // first thread consumes first point
+      if (thrx == 0) {
+        float curr = static_cast<float>(lvals[0]);
+        cuWelfordOnlineSum(curr,mu,sigma2,count);
+      }
+      ++l;
+    }
+    // at this point, lvals[l] are 32 bit aligned for all threads.
+    for (;  l+7 < n2;  l+=8*numx) {
+      for (int k = 0;  k < 8;  k+=2) {
+        float2 curr = __half22float2(*((__half2*)(lvals+l+k)));
+        cuWelfordOnlineSum(curr.x,mu,sigma2,count);
+        cuWelfordOnlineSum(curr.y,mu,sigma2,count);
+      }
+    }
+    for (;  l < n2;  ++l) {
+      float curr = static_cast<float>(lvals[l]);
+      cuWelfordOnlineSum(curr,mu,sigma2,count);
+    }
+    // intra-warp reductions
+    for (int l = 0;  l <= 4;  ++l) {
+      int srcLaneB = (threadIdx.x+(1<<l))&31;
+      float muB = WARP_SHFL(mu, srcLaneB);
+      float countB = WARP_SHFL(count, srcLaneB);
+      float sigma2B = WARP_SHFL(sigma2, srcLaneB);
+      cuChanOnlineSum(muB,sigma2B,countB,mu,sigma2,count);
+    }
+    // threadIdx.x == 0 has correct values for each warp
+    // inter-warp reductions
+    if (blockDim.y > 1) {
+      float* ubuf = (float*)buf;
+      float* ibuf = (float*)(ubuf + blockDim.y);
+      for (int offset = blockDim.y/2;  offset > 0;  offset /= 2) {
+        // upper half of warps write to shared
+        if (threadIdx.x == 0 && threadIdx.y >= offset && threadIdx.y < 2*offset) {
+          const int wrt_y = threadIdx.y - offset;
+          ubuf[2*wrt_y] = mu;
+          ubuf[2*wrt_y+1] = sigma2;
+          ibuf[wrt_y] = count;
+        }
+        __syncthreads();
+        // lower half merges
+        if (threadIdx.x == 0 && threadIdx.y < offset) {
+          float muB = ubuf[2*threadIdx.y];
+          float sigma2B = ubuf[2*threadIdx.y+1];
+          float countB = ibuf[threadIdx.y];
+          cuChanOnlineSum(muB,sigma2B,countB,mu,sigma2,count);
+        }
+        __syncthreads();
+      }
+      // threadIdx.x = 0 && threadIdx.y == 0 only thread that has correct values
+      if (threadIdx.x == 0 && threadIdx.y == 0) {
+        ubuf[0] = mu;
+        ubuf[1] = sigma2;
+      }
+      __syncthreads();
+      mu = ubuf[0];
+      sigma2 = ubuf[1]/float(n2);
+      // don't care about final value of count, we know count == n2
+    } else {
+      mu = WARP_SHFL(mu, 0);
+      sigma2 = WARP_SHFL(sigma2/float(n2), 0);
+    }
+  }
+}
+
+template<typename U> U rsqrt(U v) {
+  return U(1) / sqrt(v);
+}
+template<> float rsqrt(float v) {
+  return rsqrtf(v);
+}
+template<> double rsqrt(double v) {
+  return rsqrt(v);
+}
+
+namespace {
+// This is the un-specialized struct.  Note that we prevent instantiation of this
+// struct by putting an undefined symbol in the function body so it won't compile.
+//  template <typename T>
+//  struct SharedMemory
+//  {
+//      // Ensure that we won't compile any un-specialized types
+//      __device__ T *getPointer()
+//      {
+//          extern __device__ void error(void);
+//          error();
+//          return NULL;
+//      }
+//  };
+// https://github.com/NVIDIA/apex/issues/246
+template <typename T>
+struct SharedMemory;
+
+template <>
+struct SharedMemory <float>
+{
+    __device__ float *getPointer()
+    {
+        extern __shared__ float s_float[];
+        return s_float;
+    }
+};
+
+template <>
+struct SharedMemory <double>
+{
+    __device__ double *getPointer()
+    {
+        extern __shared__ double s_double[];
+        return s_double;
+    }
+};
+}
+
+template<typename T, typename U, typename V> __device__
+void cuApplyLayerNorm_(
+  V* __restrict__ output_vals,
+  U* __restrict__ mean,
+  U* __restrict__ invvar,
+  const T* __restrict__ vals,
+  const int n1,
+  const int n2,
+  const U epsilon,
+  const V* __restrict__ gamma,
+  const V* __restrict__ beta
+  )
+{
+  // Assumptions:
+  // 1) blockDim.x == warpSize
+  // 2) Tensors are contiguous
+  //
+  for (auto i1=blockIdx.y; i1 < n1; i1 += gridDim.y) {
+    SharedMemory<U> shared;
+    U* buf = shared.getPointer();
+    U mu,sigma2;
+    cuWelfordMuSigma2(vals,n1,n2,i1,mu,sigma2,buf);
+    const T* lvals = vals + i1*n2;
+    V* ovals = output_vals + i1*n2;
+    U c_invvar = rsqrt(sigma2 + epsilon);
+    const int numx = blockDim.x * blockDim.y;
+    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
+    if (gamma != NULL && beta != NULL) {
+      for (int i = thrx;  i < n2;  i+=numx) {
+        U curr = static_cast<U>(lvals[i]);
+        ovals[i] = gamma[i] * static_cast<V>(c_invvar * (curr - mu)) + beta[i];
+      }
+    } else {
+      for (int i = thrx;  i < n2;  i+=numx) {
+        U curr = static_cast<U>(lvals[i]);
+        ovals[i] = static_cast<V>(c_invvar * (curr - mu));
+      }
+    }
+    if (threadIdx.x == 0 && threadIdx.y == 0) {
+      mean[i1] = mu;
+      invvar[i1] = c_invvar;
+    }
+    __syncthreads();
+  }
+}
+
+template<typename T, typename U, typename V=T> __global__
+void cuApplyLayerNorm(
+  V* __restrict__ output_vals,
+  U* __restrict__ mean,
+  U* __restrict__ invvar,
+  const T* __restrict__ vals,
+  const int n1,
+  const int n2,
+  const U epsilon,
+  const V* __restrict__ gamma,
+  const V* __restrict__ beta
+  )
+{
+  cuApplyLayerNorm_<T, U, V>(output_vals, mean, invvar, vals, n1, n2, epsilon, gamma, beta);
+}
+
+
+template<typename T, typename U, typename V> __device__
+void cuLoadWriteStridedInputs(
+    const int i1_block,
+    const int thr_load_row_off,
+    const int thr_load_col_off,
+    const int i2_off,
+    const int row_stride,
+    U* warp_buf1,
+    U* warp_buf2,
+    const T* input,
+    const V* dout,
+    const int i1_end,
+    const int n2,
+    const U* __restrict__ mean,
+    const U* __restrict__ invvar
+    )
+{
+  int i1 = i1_block+thr_load_row_off;
+  if (i1 < i1_end) {
+    U curr_mean = mean[i1];
+    U curr_invvar = invvar[i1];
+    for (int k = 0;  k < blockDim.y;  ++k) {
+      int i2 = i2_off + k;
+      int load_idx = i1*n2+i2;
+      int write_idx = thr_load_row_off*row_stride+thr_load_col_off+k;
+      if (i2<n2) {
+        U curr_input = static_cast<U>(input[load_idx]);
+        U curr_dout = static_cast<U>(dout[load_idx]);
+        warp_buf1[write_idx] = curr_dout;
+        warp_buf2[write_idx] = curr_dout * (curr_input - curr_mean) * curr_invvar;
+      } else {
+        warp_buf1[write_idx] = U(0);
+        warp_buf2[write_idx] = U(0);
+      }
+    }
+  } else {
+    for (int k = 0;  k < blockDim.y;  ++k) {
+      int write_idx = thr_load_row_off*row_stride+thr_load_col_off+k;
+      warp_buf1[write_idx] = U(0);
+      warp_buf2[write_idx] = U(0);
+    }
+  }
+}
+
+template<typename T, typename U, typename V> __device__
+void cuLoadAddStridedInputs(
+    const int i1_block,
+    const int thr_load_row_off,
+    const int thr_load_col_off,
+    const int i2_off,
+    const int row_stride,
+    U* warp_buf1,
+    U* warp_buf2,
+    const T* input,
+    const V* dout,
+    const int i1_end,
+    const int n2,
+    const U* __restrict__ mean,
+    const U* __restrict__ invvar
+    )
+{
+  int i1 = i1_block+thr_load_row_off;
+  if (i1 < i1_end) {
+    U curr_mean = mean[i1];
+    U curr_invvar = invvar[i1];
+    for (int k = 0;  k < blockDim.y;  ++k) {
+      int i2 = i2_off + k;
+      int load_idx = i1*n2+i2;
+      int write_idx = thr_load_row_off*row_stride+thr_load_col_off+k;
+      if (i2<n2) {
+        U curr_input = static_cast<U>(input[load_idx]);
+        U curr_dout = static_cast<U>(dout[load_idx]);
+        warp_buf1[write_idx] += curr_dout;
+        warp_buf2[write_idx] += curr_dout * (curr_input - curr_mean) * curr_invvar;
+      }
+    }
+  }
+}
+
+template<typename T, typename U, typename V> __global__
+void cuComputePartGradGammaBeta(
+    const V* __restrict__ dout,
+    const T* __restrict__ input,
+    const int n1,
+    const int n2,
+    const U* __restrict__ mean,
+    const U* __restrict__ invvar,
+    U epsilon,
+    U* part_grad_gamma,
+    U* part_grad_beta)
+{
+    const int numsegs_n1 = (n1+blockDim.y*blockDim.y-1) / (blockDim.y*blockDim.y);
+    const int segs_per_block = (numsegs_n1 + gridDim.y - 1) / gridDim.y;
+    const int i1_beg = blockIdx.y * segs_per_block * blockDim.y*blockDim.y;
+    const int i1_beg_plus_one = (blockIdx.y+1) * segs_per_block * blockDim.y*blockDim.y;
+    const int i1_end = i1_beg_plus_one < n1 ? i1_beg_plus_one : n1;
+    const int row_stride = blockDim.x+1;
+    const int thr_load_col_off = (threadIdx.x*blockDim.y)&(blockDim.x-1);
+    const int thr_load_row_off = (threadIdx.x*blockDim.y)/blockDim.x + threadIdx.y*blockDim.y;
+    const int i2_off = blockIdx.x * blockDim.x + thr_load_col_off;
+    SharedMemory<U> shared;
+    U* buf = shared.getPointer(); // buf has at least blockDim.x * blockDim.y * blockDim.y + (blockDim.y - 1)*(blockDim.x/blockDim.y) elements
+    U* warp_buf1 = (U*)buf;
+    U* warp_buf2 = warp_buf1 + blockDim.y * blockDim.y * row_stride;
+    // compute partial sums from strided inputs
+    // do this to increase number of loads in flight
+    cuLoadWriteStridedInputs(i1_beg,thr_load_row_off,thr_load_col_off,i2_off,row_stride,warp_buf1,warp_buf2,input,dout,i1_end,n2,mean,invvar);
+    for (int i1_block = i1_beg+blockDim.y*blockDim.y;  i1_block < i1_end;  i1_block+=blockDim.y*blockDim.y) {
+      cuLoadAddStridedInputs(i1_block,thr_load_row_off,thr_load_col_off,i2_off,row_stride,warp_buf1,warp_buf2,input,dout,i1_end,n2,mean,invvar);
+    }
+    __syncthreads();
+    // inter-warp reductions
+    // sum within each warp
+    U acc1 = U(0);
+    U acc2 = U(0);
+    for (int k = 0;  k < blockDim.y;  ++k) {
+      int row1 = threadIdx.y + k*blockDim.y;
+      int idx1 = row1*row_stride + threadIdx.x;
+      acc1 += warp_buf1[idx1];
+      acc2 += warp_buf2[idx1];
+    }
+    warp_buf1[threadIdx.y*row_stride+threadIdx.x] = acc1;
+    warp_buf2[threadIdx.y*row_stride+threadIdx.x] = acc2;
+    __syncthreads();
+    // sum all warps
+    for (int offset = blockDim.y/2;  offset > 1;  offset /= 2) {
+      if (threadIdx.y < offset) {
+        int row1 = threadIdx.y;
+        int row2 = threadIdx.y + offset;
+        int idx1 = row1*row_stride + threadIdx.x;
+        int idx2 = row2*row_stride + threadIdx.x;
+        warp_buf1[idx1] += warp_buf1[idx2];
+        warp_buf2[idx1] += warp_buf2[idx2];
+      }
+      __syncthreads();
+    }
+    int i2 = blockIdx.x * blockDim.x + threadIdx.x;
+    if (threadIdx.y == 0 && i2 < n2) {
+      int row1 = threadIdx.y;
+      int row2 = threadIdx.y + 1;
+      int idx1 = row1*row_stride + threadIdx.x;
+      int idx2 = row2*row_stride + threadIdx.x;
+      part_grad_beta[blockIdx.y*n2+i2] = warp_buf1[idx1] + warp_buf1[idx2];
+      part_grad_gamma[blockIdx.y*n2+i2] = warp_buf2[idx1] + warp_buf2[idx2];
+    }
+}
+
+template<typename U, typename V> __global__
+void cuComputeGradGammaBeta(
+    const U* part_grad_gamma,
+    const U* part_grad_beta,
+    const int part_size,
+    const int n1,
+    const int n2,
+    V* grad_gamma,
+    V* grad_beta)
+{
+    // sum partial gradients for gamma and beta
+    SharedMemory<U> shared;
+    U* buf = shared.getPointer();
+    int i2 = blockIdx.x * blockDim.x + threadIdx.x;
+    if (i2 < n2) {
+      // each warp does sequential reductions until reduced part_size is num_warps
+      int num_warp_reductions = part_size / blockDim.y;
+      U sum_gamma = U(0);
+      U sum_beta = U(0);
+      const U* part_grad_gamma_ptr = part_grad_gamma + threadIdx.y * num_warp_reductions * n2 + i2;
+      const U* part_grad_beta_ptr = part_grad_beta + threadIdx.y * num_warp_reductions * n2 + i2;
+      for (int warp_offset = 0;  warp_offset < num_warp_reductions;  ++warp_offset) {
+        sum_gamma += part_grad_gamma_ptr[warp_offset*n2];
+        sum_beta += part_grad_beta_ptr[warp_offset*n2];
+      }
+      // inter-warp reductions
+      const int nbsize3 = blockDim.x * blockDim.y / 2;
+      for (int offset = blockDim.y/2;  offset >= 1;  offset /= 2) {
+        // top half write to shared memory
+        if (threadIdx.y >= offset && threadIdx.y < 2*offset) {
+          const int write_idx = (threadIdx.y - offset) * blockDim.x + threadIdx.x;
+          buf[write_idx] = sum_gamma;
+          buf[write_idx+nbsize3] = sum_beta;
+        }
+        __syncthreads();
+        // bottom half sums
+        if (threadIdx.y < offset) {
+          const int read_idx = threadIdx.y * blockDim.x + threadIdx.x;
+          sum_gamma += buf[read_idx];
+          sum_beta += buf[read_idx+nbsize3];
+        }
+        __syncthreads();
+      }
+      // write out fully summed gradients
+      if (threadIdx.y == 0) {
+        grad_gamma[i2] = sum_gamma;
+        grad_beta[i2] = sum_beta;
+      }
+    }
+}
+
+template<typename T, typename U, typename V> __global__
+void cuComputeGradInput(
+    const V* __restrict__ dout,
+    const T* __restrict__ input,
+    const int n1,
+    const int n2,
+    const U* __restrict__ mean,
+    const U* __restrict__ invvar,
+    U epsilon,
+    const V* gamma,
+    T* grad_input)
+{
+  for (auto i1=blockIdx.y; i1 < n1; i1 += gridDim.y) {
+    U sum_loss1 = U(0);
+    U sum_loss2 = U(0);
+    const U c_mean = mean[i1];
+    const U c_invvar = invvar[i1];
+    const T* k_input = input + i1*n2;
+    const V* k_dout = dout + i1*n2;
+    const int numx = blockDim.x * blockDim.y;
+    const int thrx = threadIdx.x + threadIdx.y * blockDim.x;
+    if (gamma != NULL) {
+      int l = 4*thrx;
+      for (;  l+3 < n2;  l+=4*numx) {
+        for (int k = 0;  k < 4;  ++k) {
+          const U c_h = static_cast<U>(k_input[l+k]);
+          const U c_loss = static_cast<U>(k_dout[l+k]);
+          sum_loss1 += c_loss * gamma[l+k];
+          sum_loss2 += c_loss * gamma[l+k] * (c_h - c_mean) * c_invvar;
+        }
+      }
+      for (;  l < n2;  ++l) {
+        const U c_h = static_cast<U>(k_input[l]);
+        const U c_loss = static_cast<U>(k_dout[l]);
+        sum_loss1 += c_loss * gamma[l];
+        sum_loss2 += c_loss * gamma[l] * (c_h - c_mean) * c_invvar;
+      }
+    } else {
+      int l = 4*thrx;
+      for (;  l+3 < n2;  l+=4*numx) {
+        for (int k = 0;  k < 4;  ++k) {
+          const U c_h = static_cast<U>(k_input[l+k]);
+          const U c_loss = static_cast<U>(k_dout[l+k]);
+          sum_loss1 += c_loss;
+          sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;
+        }
+      }
+      for (;  l < n2;  ++l) {
+        const U c_h = static_cast<U>(k_input[l]);
+        const U c_loss = static_cast<U>(k_dout[l]);
+        sum_loss1 += c_loss;
+        sum_loss2 += c_loss * (c_h - c_mean) * c_invvar;
+      }
+    }
+    // intra-warp reductions
+    for (int mask = blockDim.x/2;  mask > 0;  mask /= 2) {
+      sum_loss1 += WARP_SHFL_XOR(sum_loss1, mask);
+      sum_loss2 += WARP_SHFL_XOR(sum_loss2, mask);
+    }
+    // inter-warp reductions
+    if (blockDim.y > 1) {
+      SharedMemory<U> shared;
+      U* buf = shared.getPointer();
+      for (int offset = blockDim.y/2;  offset > 0;  offset /= 2) {
+        // upper half of warps write to shared
+        if (threadIdx.y >= offset && threadIdx.y < 2*offset) {
+          const int wrt_i = (threadIdx.y - offset) * blockDim.x + threadIdx.x;
+          buf[2*wrt_i] = sum_loss1;
+          buf[2*wrt_i+1] = sum_loss2;
+        }
+        __syncthreads();
+        // lower half merges
+        if (threadIdx.y < offset) {
+          const int read_i = threadIdx.y * blockDim.x + threadIdx.x;
+          sum_loss1 += buf[2*read_i];
+          sum_loss2 += buf[2*read_i+1];
+        }
+        __syncthreads();
+      }
+      if (threadIdx.y == 0) {
+        buf[2*threadIdx.x] = sum_loss1;
+        buf[2*threadIdx.x+1] = sum_loss2;
+      }
+      __syncthreads();
+      if (threadIdx.y !=0) {
+        sum_loss1 = buf[2*threadIdx.x];
+        sum_loss2 = buf[2*threadIdx.x+1];
+      }
+    }
+    // all threads now have the two sums over l
+    U fH = (U)n2;
+    U term1 = (U(1) / fH) * c_invvar;
+    T* k_grad_input = grad_input + i1*n2;
+    if (gamma != NULL) {
+      for (int l = thrx;  l < n2;  l+=numx) {
+        const U c_h = static_cast<U>(k_input[l]);
+        const U c_loss = static_cast<U>(k_dout[l]);
+        U f_grad_input = fH * c_loss * gamma[l];
+        f_grad_input -= sum_loss1;
+        f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;
+        f_grad_input *= term1;
+        k_grad_input[l] = static_cast<T>(f_grad_input);
+      }
+    } else {
+      for (int l = thrx;  l < n2;  l+=numx) {
+        const U c_h = static_cast<U>(k_input[l]);
+        const U c_loss = static_cast<U>(k_dout[l]);
+        U f_grad_input = fH * c_loss;
+        f_grad_input -= sum_loss1;
+        f_grad_input -= (c_h - c_mean) * c_invvar * sum_loss2;
+        f_grad_input *= term1;
+        k_grad_input[l] = static_cast<T>(f_grad_input);
+      }
+    }
+    // prevent race where buf is written again before reads are done
+    __syncthreads();
+  }
+}
+
+template<typename T, typename U, typename V=T>
+void HostApplyLayerNorm(
+    V* output,
+    U* mean,
+    U* invvar,
+    const T* input,
+    int n1,
+    int n2,
+    double epsilon,
+    const V* gamma,
+    const V* beta
+    )
+{
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+    const dim3 threads(32,4,1);
+    const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];
+    const dim3 blocks(1, std::min((uint64_t)n1, maxGridY), 1);
+    int nshared =
+        threads.y > 1 ?
+	    threads.y*sizeof(U)+(threads.y/2)*sizeof(U) :
+	    0;
+    cuApplyLayerNorm<<<blocks, threads, nshared, stream>>>(
+      output, mean, invvar, input, n1, n2, U(epsilon), gamma, beta);
+}
+
+void cuda_layer_norm(
+    at::Tensor* output,
+    at::Tensor* mean,
+    at::Tensor* invvar,
+    at::Tensor* input,
+    int n1,
+    int n2,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    at::Tensor* gamma,
+    at::Tensor* beta,
+    double epsilon)
+{
+    using namespace at;
+    DISPATCH_DOUBLE_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(
+        input->scalar_type(), output->scalar_type(), "layer_norm_cuda_kernel",
+        using accscalar_t = at::acc_type<scalar_t_in, true>;
+        HostApplyLayerNorm<scalar_t_in, accscalar_t, scalar_t_out>(
+          output->DATA_PTR<scalar_t_out>(),
+	      mean->DATA_PTR<accscalar_t>(),
+          invvar->DATA_PTR<accscalar_t>(),
+          input->DATA_PTR<scalar_t_in>(),
+          n1,n2,
+          epsilon,
+          gamma != NULL ? gamma->DATA_PTR<scalar_t_out>() : NULL,
+          beta != NULL ? beta->DATA_PTR<scalar_t_out>() : NULL);
+      )
+}
+
+template<typename T, typename U=float, typename V=T>
+void HostLayerNormGradient(
+    const V* dout,
+    const U* mean,
+    const U* invvar,
+    at::Tensor* input,
+    int n1,
+    int n2,
+    const V* gamma,
+    const V* beta,
+    double epsilon,
+    T* grad_input,
+    V* grad_gamma,
+    V* grad_beta
+    )
+{
+    auto stream = at::cuda::getCurrentCUDAStream().stream();
+
+    if (gamma != NULL && beta != NULL) {
+      // compute grad_gamma(j) and grad_beta(j)
+      const int part_size = 16;
+      const dim3 threads2(32,4,1);
+      const dim3 blocks2((n2+threads2.x-1)/threads2.x,part_size,1);
+      const int nshared2_a = 2 * sizeof(U) * threads2.y * threads2.y * (threads2.x + 1);
+      const int nshared2_b = threads2.x * threads2.y * sizeof(U);
+      const int nshared2 = nshared2_a > nshared2_b ? nshared2_a : nshared2_b;
+      // note (mkozuki): I can hard code part_grad_gamma's dtype as float given that
+      // the `cuda_layer_norm_gradient` doesn't support double.
+      const auto part_grad_dtype =
+        (input->scalar_type() == at::ScalarType::Half || input->scalar_type() == at::ScalarType::BFloat16) ?
+        at::ScalarType::Float :
+        input->scalar_type();
+      at::Tensor part_grad_gamma = at::empty({part_size,n2}, input->options().dtype(part_grad_dtype));
+      at::Tensor part_grad_beta = at::empty_like(part_grad_gamma);
+      cuComputePartGradGammaBeta<<<blocks2, threads2, nshared2, stream>>>(
+		      dout,
+		      input->DATA_PTR<T>(),
+		      n1,n2,
+		      mean,
+		      invvar,
+		      U(epsilon),
+		      part_grad_gamma.DATA_PTR<U>(),
+		      part_grad_beta.DATA_PTR<U>());
+
+      const dim3 threads3(32,8,1);
+      const dim3 blocks3((n2+threads2.x-1)/threads2.x,1,1);
+      const int nshared3 = threads3.x * threads3.y * sizeof(U);
+      cuComputeGradGammaBeta<<<blocks3, threads3, nshared3, stream>>>(
+		      part_grad_gamma.DATA_PTR<U>(),
+		      part_grad_beta.DATA_PTR<U>(),
+		      part_size,
+		      n1,n2,
+		      grad_gamma,
+		      grad_beta);
+    }
+
+    // compute grad_input
+    const uint64_t maxGridY = at::cuda::getCurrentDeviceProperties()->maxGridSize[1];
+    const dim3 blocks1(1, std::min((uint64_t)n1, maxGridY), 1);
+    const dim3 threads1(32,4,1);
+    int nshared =
+	    threads1.y > 1 ?
+	    threads1.y*threads1.x*sizeof(U) :
+	    0;
+    cuComputeGradInput<<<blocks1, threads1, nshared, stream>>>(
+            dout,
+            input->DATA_PTR<T>(),
+            n1,n2,
+            mean,
+            invvar,
+            U(epsilon),
+            gamma,
+            grad_input);
+}
+
+void cuda_layer_norm_gradient(
+    at::Tensor* dout,
+    at::Tensor* mean,
+    at::Tensor* invvar,
+    at::Tensor* input,
+    int n1,
+    int n2,
+    #ifdef VERSION_GE_1_1
+    at::IntArrayRef normalized_shape,
+    #else
+    at::IntList normalized_shape,
+    #endif
+    at::Tensor* gamma,
+    at::Tensor* beta,
+    double epsilon,
+    at::Tensor* grad_input,
+    at::Tensor* grad_gamma,
+    at::Tensor* grad_beta)
+{
+    using namespace at;
+    // we can do away with `accscalar_t` as there're only three dtypes: fp32, fp16, bf16
+    DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(
+      input->scalar_type(), gamma == NULL ? input->scalar_type() :  gamma->scalar_type(), "cuComputeGradInput",
+      using accscalar_t = at::acc_type<scalar_t_in, true>;
+      HostLayerNormGradient(
+        dout->DATA_PTR<scalar_t_out>(),
+        mean->DATA_PTR<accscalar_t>(),
+        invvar->DATA_PTR<accscalar_t>(),
+        input,
+        n1,n2,
+            // TMJ pass NULL argument for gamma, beta, grad_gamma and grad_beta
+            // if gamma Tensor is NULL on input.
+        gamma != NULL ? gamma->DATA_PTR<scalar_t_out>() : NULL,
+        gamma != NULL ? beta->DATA_PTR<scalar_t_out>() : NULL,
+        epsilon,
+        grad_input->DATA_PTR<scalar_t_in>(),
+        gamma != NULL ? grad_gamma->DATA_PTR<scalar_t_out>() : NULL,
+        gamma != NULL ? grad_beta->DATA_PTR<scalar_t_out>() : NULL);
+    )
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax.cpp
new file mode 100644
index 0000000000..6e5d355649
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax.cpp
@@ -0,0 +1,97 @@
+/* coding=utf-8
+ * Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+#include <vector>
+
+namespace multihead_attn {
+namespace fused_softmax {
+namespace scaled_masked_softmax {
+
+torch::Tensor fwd_cuda(
+    torch::Tensor const& input, 
+    torch::Tensor const& mask,
+    float scale_factor);
+
+torch::Tensor bwd_cuda(
+    torch::Tensor const& output_grads, 
+    torch::Tensor const& softmax_results,
+    float scale_factor);
+
+int get_batch_per_block_cuda(
+    int query_seq_len,
+    int key_seq_len,
+    int batches,
+    int attn_heads);
+
+torch::Tensor fwd(
+    torch::Tensor const& input,
+    torch::Tensor const& mask,
+    float scale_factor) {
+  AT_ASSERTM(input.dim() == 4, "expected 4D tensor");
+  AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) ||
+	     (input.scalar_type() == at::ScalarType::BFloat16), 
+      "Only fp16 and bf16 are supported");
+  AT_ASSERTM(mask.dim() == 4, "expected 4D tensor");
+
+  return fwd_cuda(input, mask, scale_factor);
+}
+
+torch::Tensor bwd(
+    torch::Tensor const& output_grads, 
+    torch::Tensor const& softmax_results,
+    float scale_factor) {
+
+  AT_ASSERTM(output_grads.dim() == 4, "expected 3D tensor");
+  AT_ASSERTM(softmax_results.dim() == 4, "expected 3D tensor");
+
+  AT_ASSERTM((output_grads.scalar_type() == at::ScalarType::Half) ||
+	     (output_grads.scalar_type() == at::ScalarType::BFloat16), 
+      "Only fp16 and bf16 are supported");
+  AT_ASSERTM((softmax_results.scalar_type() == at::ScalarType::Half) ||
+	     (softmax_results.scalar_type() == at::ScalarType::BFloat16), 
+      "Only fp16 and bf16 are supported");
+
+  return bwd_cuda(output_grads, softmax_results, scale_factor);
+}
+
+int get_batch_per_block(
+    int query_seq_len,
+    int key_seq_len,
+    int batches,
+    int attn_heads) {
+    return get_batch_per_block_cuda(query_seq_len, key_seq_len, batches, attn_heads);
+}
+
+} // end namespace scaled_masked_softmax
+} // end namespace fused_softmax
+} // end namespace multihead_attn
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", 
+        &multihead_attn::fused_softmax::scaled_masked_softmax::fwd, 
+	"Self Multihead Attention scaled, time masked softmax -- Forward.");
+
+  m.def("backward",
+        &multihead_attn::fused_softmax::scaled_masked_softmax::bwd,
+	"Self Multihead Attention scaled, time masked softmax -- Backward.");
+
+  m.def("get_batch_per_block",
+        &multihead_attn::fused_softmax::scaled_masked_softmax::get_batch_per_block,
+        "Return Batch per block size."
+  );
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax.h
new file mode 100644
index 0000000000..78a29cf3ba
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax.h
@@ -0,0 +1,505 @@
+/* coding=utf-8
+ * Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#pragma once
+
+#include <assert.h>
+#include <cuda_fp16.h>
+#include <cfloat>
+#include <limits>
+#include <stdint.h>
+#include <cuda_fp16.h>
+#include <c10/macros/Macros.h>
+
+namespace {
+
+template <typename Datatype, int ELEMENTS_PER_LDG>
+__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);
+
+template <>
+__device__ __inline__ void copy_vector<c10::BFloat16, 1>(c10::BFloat16 *dst, const c10::BFloat16 *src) { *dst = *src; }
+
+template <>
+__device__ __inline__ void copy_vector<c10::BFloat16, 4>(c10::BFloat16 *dst, const c10::BFloat16 *src) { *((float2*) dst) = *((float2*) src); }
+
+template <>
+__device__ __inline__ void copy_vector<c10::Half, 1>(c10::Half *dst, const c10::Half *src) { *dst = *src; }
+
+template <>
+__device__ __inline__ void copy_vector<c10::Half, 4>(c10::Half *dst, const c10::Half *src) { *((float2*) dst) = *((float2*) src); }
+
+template <>
+__device__ __inline__ void copy_vector<uint8_t, 1>(uint8_t *dst, const uint8_t *src) { *dst = *src; }
+
+template <>
+__device__ __inline__ void copy_vector<uint8_t, 4>(uint8_t *dst, const uint8_t *src) {*((half2*) dst) = *((half2*) src); }
+
+int log2_ceil(int value) {
+    int log2_value = 0;
+    while ((1 << log2_value) < value) ++log2_value;
+    return log2_value;
+}
+
+template<typename T>
+struct Add {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a + b;
+  }
+};
+
+template<typename T>
+struct Max {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a < b ? b : a;
+  }
+};
+
+template <typename T>
+__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff)
+{
+#if CUDA_VERSION >= 9000
+    return __shfl_xor_sync(mask, value, laneMask, width);
+#else
+    return __shfl_xor(value, laneMask, width);
+#endif
+}
+
+template <typename acc_t, int WARP_BATCH, int WARP_SIZE, template<typename> class ReduceOp>
+__device__ __forceinline__ void warp_reduce(acc_t* sum) {
+    ReduceOp<acc_t> r;
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;  i < WARP_BATCH;  ++i) {
+            acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);
+            sum[i] = r(sum[i], b);
+        }
+    }
+}
+
+/*
+ * Extended softmax (from native aten pytorch) with following additional features
+ * 1) input scaling
+ * 2) Explicit masking
+ */	
+template <typename input_t, typename output_t, typename acc_t, int log2_elements>
+__global__ void scaled_masked_softmax_warp_forward(
+    output_t *dst, 
+    const input_t *src,
+    const uint8_t *mask, 
+    const acc_t scale, 
+    int micro_batch_size, 
+    int element_count,
+    int pad_batches) 
+{
+    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and 
+    // warp_size of method warp_softmax_forward_kernel.
+    constexpr int next_power_of_two = 1 << log2_elements;
+    constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+    constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
+    constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
+    constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;
+
+    // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )
+    // gridDim/blockIdx = (seq_len, attn_heads, batches) 
+    int first_batch = (blockDim.y * (blockIdx.x + gridDim.x * (blockIdx.y + gridDim.y * blockIdx.z))+ threadIdx.y) * WARP_BATCH;
+    int pad_first_batch = 0;
+    if (pad_batches != 1) { // bert style
+        pad_first_batch = (blockDim.y * (blockIdx.x + gridDim.x * blockIdx.z) + threadIdx.y) * WARP_BATCH;
+    } else { // gpt2 style
+        pad_first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+    }
+
+    // micro_batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = micro_batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+
+    src += first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;
+    dst += first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;
+    mask += pad_first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;
+
+    // load data from global memory
+    acc_t elements[WARP_BATCH][WARP_ITERATIONS];
+    input_t temp_data[ELEMENTS_PER_LDG_STG];
+    uint8_t temp_mask[ELEMENTS_PER_LDG_STG];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  it+=ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+
+            if (element_index < batch_element_count) {
+                int itr_idx = i*element_count+it*WARP_SIZE;
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(temp_data, src + itr_idx);
+                copy_vector<uint8_t, ELEMENTS_PER_LDG_STG>(temp_mask, mask + itr_idx);
+
+                #pragma unroll
+                  for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                      if (temp_mask[element] != 1) {
+                          elements[i][it + element] = (acc_t)temp_data[element] * scale;
+                      } else {
+                          elements[i][it + element] = -10000.0;
+                      }
+                  }
+            } else {
+                #pragma unroll
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    elements[i][it + element] = -std::numeric_limits<acc_t>::infinity();
+                }
+            }
+        }
+    }
+
+    // compute max_value
+    acc_t max_value[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        max_value[i] = elements[i][0];
+        #pragma unroll
+        for (int it = 1;  it < WARP_ITERATIONS;  ++it) {
+            max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
+        }
+    }
+    warp_reduce<acc_t, WARP_BATCH, WARP_SIZE, Max>(max_value);
+
+    acc_t sum[WARP_BATCH] { 0.0f };
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  ++it) {
+            elements[i][it] = std::exp((elements[i][it] - max_value[i]));
+            sum[i] += elements[i][it];
+        }
+    }
+    warp_reduce<acc_t, WARP_BATCH, WARP_SIZE, Add>(sum);
+
+    // store result
+    output_t out[ELEMENTS_PER_LDG_STG];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  it+=ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                #pragma unroll
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    out[element] = elements[i][it + element] / sum[i];
+                }
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(dst + i * element_count + it * WARP_SIZE, out);  
+            } else {
+                break;
+            } 
+        }
+    }
+}
+
+template <typename input_t, typename output_t, typename acc_t, int log2_elements>
+__global__ void scaled_masked_softmax_warp_backward(
+    output_t *gradInput, 
+    input_t *grad, 
+    const input_t *output,
+    acc_t scale, 
+    int micro_batch_size, 
+    int element_count)
+{
+    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and 
+    // warp_size of method warp_softmax_backward_kernel.
+    constexpr int next_power_of_two = 1 << log2_elements;
+    constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+    constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
+    constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
+    constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;
+
+    // blockDim/threadIdx = (WARP_SIZE, WARPS_PER_BLOCK, )
+    // gridDim/blockIdx = (seq_len, attn_heads, batches) 
+    int first_batch = (blockDim.y * blockIdx.x + threadIdx.y) * WARP_BATCH;
+    
+    // micro_batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = micro_batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+
+    // the first element to process by the current thread
+    int thread_offset = first_batch * element_count + ELEMENTS_PER_LDG_STG * local_idx;
+    grad += thread_offset;
+    output += thread_offset;
+    gradInput += thread_offset;
+
+    // load data from global memory
+    acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS] { 0.0f };
+    acc_t output_reg[WARP_BATCH][WARP_ITERATIONS] { 0.0f };
+    input_t temp_grad[ELEMENTS_PER_LDG_STG];
+    input_t temp_output[ELEMENTS_PER_LDG_STG];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : element_count;
+
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  it+=ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < batch_element_count) {
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(temp_grad, grad + i * element_count + it * WARP_SIZE);
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(temp_output, output + i * element_count + it * WARP_SIZE);
+
+                #pragma unroll
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    output_reg[i][it + element] = (acc_t)temp_output[element];
+                }
+                #pragma unroll
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    grad_reg[i][it + element] = (acc_t)temp_grad[element] * output_reg[i][it + element];
+                }
+            } 
+        }
+    }
+   
+    acc_t sum[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        sum[i] = grad_reg[i][0];
+        #pragma unroll
+        for (int it = 1;  it < WARP_ITERATIONS;  ++it) {
+            sum[i] += grad_reg[i][it];
+        }
+    }
+    warp_reduce<acc_t, WARP_BATCH, WARP_SIZE, Add>(sum);
+
+    // store result
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  it+=ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                // compute gradients
+                output_t out[ELEMENTS_PER_LDG_STG];
+                #pragma unroll
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    out[element] = (output_t)(scale * (grad_reg[i][it + element] - output_reg[i][it + element] * sum[i]));
+                }
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(gradInput + i * element_count + it * WARP_SIZE, out);
+            } 
+        }
+    }
+}
+} // end of anonymous namespace
+
+int get_batch_per_block(int query_seq_len, int key_seq_len, int batches, int attn_heads){
+    int log2_elements = log2_ceil(key_seq_len);
+    const int next_power_of_two = 1 << log2_elements;
+
+    int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+    int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+
+    constexpr int threads_per_block = 128;
+    int warps_per_block = (threads_per_block / warp_size);
+    int batches_per_block = warps_per_block * batches_per_warp;
+
+    return batches_per_block;
+}
+
+template<typename input_t, typename output_t, typename acc_t>
+void dispatch_scaled_masked_softmax_forward(
+    output_t *dst, 
+    const input_t *src, 
+    const uint8_t *mask,
+    const input_t scale, 
+    int query_seq_len, 
+    int key_seq_len, 
+    int batches,
+    int attn_heads,
+    int pad_batches)
+{
+    TORCH_INTERNAL_ASSERT(key_seq_len >= 0 && key_seq_len <= 2048 );
+    if (key_seq_len == 0) {
+        return;
+    } else {
+        int log2_elements = log2_ceil(key_seq_len);
+        const int next_power_of_two = 1 << log2_elements;
+        int batch_count = batches * attn_heads * query_seq_len;
+
+        // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_forward.
+        int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+
+        // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_forward.
+        int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+
+        int warps_per_block = (threads_per_block / warp_size);
+        int batches_per_block = warps_per_block * batches_per_warp;
+        TORCH_INTERNAL_ASSERT(query_seq_len%batches_per_block == 0);
+        dim3 blocks(query_seq_len/batches_per_block, attn_heads, batches);
+        dim3 threads(warp_size, warps_per_block, 1);
+        // Launch code would be more elegant if C++ supported FOR CONSTEXPR
+        switch (log2_elements) {
+            case 0: // 1
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 0>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 1: // 2
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 1>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 2: // 4
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 2>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 3: // 8
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 3>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 4: // 16
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 4>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 5: // 32
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 5>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 6: // 64
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 6>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 7: // 128
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 7>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 8: // 256
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 8>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 9: // 512
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 9>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 10: // 1024
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 10>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            case 11: // 2048
+                scaled_masked_softmax_warp_forward<input_t, output_t, acc_t, 11>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, mask, scale, batch_count, key_seq_len, pad_batches);
+                break;
+            default:
+                break;
+        }
+    }
+}
+
+template<typename input_t, typename output_t, typename acc_t>
+void dispatch_scaled_masked_softmax_backward(
+    output_t *grad_input, 
+    input_t *grad, 
+    const input_t *output, 
+    const acc_t scale, 
+    int query_seq_len, 
+    int key_seq_len, 
+    int batches,
+    int attn_heads)
+{
+    TORCH_INTERNAL_ASSERT( key_seq_len >= 0 && key_seq_len <= 2048 );
+    if (key_seq_len == 0) {
+       return;
+    } else {
+        int log2_elements = log2_ceil(key_seq_len);
+        const int next_power_of_two = 1 << log2_elements;
+        int batch_count = batches *  attn_heads * query_seq_len;
+
+        // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_backward.
+        int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+
+        // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_backward.
+        int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+
+        int warps_per_block = (threads_per_block / warp_size);
+        int batches_per_block = warps_per_block * batches_per_warp;
+        int blocks = batch_count/batches_per_block;
+        dim3 threads(warp_size, warps_per_block, 1);
+        // Launch code would be more elegant if C++ supported FOR CONSTEXPR
+        switch (log2_elements) {
+            case 0: // 1
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 0>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 1: // 2
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 1>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 2: // 4
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 2>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 3: // 8
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 3>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 4: // 16
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 4>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 5: // 32
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 5>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 6: // 64
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 6>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 7: // 128
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 7>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 8: // 256
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 8>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 9: // 512
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 9>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 10: // 1024
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 10>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            case 11: // 2048
+                scaled_masked_softmax_warp_backward<input_t, output_t, acc_t, 11>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, key_seq_len);
+                break;
+            default:
+                break;
+        }
+    }
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax_cuda.cu
new file mode 100644
index 0000000000..12a364e44d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_masked_softmax_cuda.cu
@@ -0,0 +1,117 @@
+/* coding=utf-8
+ * Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#include <ATen/ATen.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+#include "scaled_masked_softmax.h"
+#include "type_shim.h"
+
+namespace multihead_attn {
+namespace fused_softmax {
+namespace scaled_masked_softmax {
+
+int get_batch_per_block_cuda(int query_seq_len, int key_seq_len, int batches, int attn_heads){
+    return get_batch_per_block(query_seq_len, key_seq_len, batches, attn_heads);
+}
+
+
+torch::Tensor fwd_cuda(
+    torch::Tensor const& input,
+    torch::Tensor const& mask,
+    float scale_factor)
+{
+  // input is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]
+  const int batches = input.size(0);
+  const int pad_batches = mask.size(0);
+  const int attn_heads = input.size(1);
+  const int query_seq_len = input.size(2);
+  const int key_seq_len = input.size(3);
+  TORCH_INTERNAL_ASSERT(key_seq_len <= 2048);
+  TORCH_INTERNAL_ASSERT(query_seq_len > 1);
+  TORCH_INTERNAL_ASSERT(pad_batches == 1 || pad_batches == batches);
+  TORCH_INTERNAL_ASSERT(mask.size(1) == 1);
+  TORCH_INTERNAL_ASSERT(mask.size(2) == query_seq_len);
+  TORCH_INTERNAL_ASSERT(mask.size(3) == key_seq_len);
+
+  // Output 
+  auto act_options = input.options().requires_grad(false);
+  torch::Tensor softmax_results = 
+      torch::empty({batches, attn_heads, query_seq_len, key_seq_len}, act_options);
+
+  // Softmax Intermediate Result Ptr
+  void* input_ptr = static_cast<void*>(input.data_ptr());
+  void* mask_ptr = static_cast<void*>(mask.data_ptr());
+  void* softmax_results_ptr = static_cast<void*>(softmax_results.data_ptr());
+
+  DISPATCH_HALF_AND_BFLOAT(
+      input.scalar_type(),
+      "dispatch_scaled_masked_softmax_forward",
+      dispatch_scaled_masked_softmax_forward<scalar_t, scalar_t, float>(
+          reinterpret_cast<scalar_t*>(softmax_results_ptr),
+	  reinterpret_cast<const scalar_t*>(input_ptr),
+	  reinterpret_cast<const uint8_t*>(mask_ptr),
+	  scale_factor,
+	  query_seq_len,
+	  key_seq_len,
+	  batches,
+	  attn_heads,
+	  pad_batches);
+      );
+  return softmax_results;
+}
+
+torch::Tensor bwd_cuda(
+    torch::Tensor const& output_grads_, 
+    torch::Tensor const& softmax_results_, 
+    float scale_factor)  {
+	
+  auto output_grads = output_grads_.contiguous();
+  auto softmax_results = softmax_results_.contiguous();
+
+  //output grads is a 4d tensor with dimensions [batches, attn_heads, seq_len, seq_len]
+  const int batches = output_grads.size(0);
+  const int attn_heads = output_grads.size(1);
+  const int query_seq_len = output_grads.size(2);
+  const int key_seq_len = output_grads.size(3);
+
+  void* output_grads_ptr = static_cast<void*>(output_grads.data_ptr());
+
+  //Softmax Grad
+  DISPATCH_HALF_AND_BFLOAT(
+      output_grads_.scalar_type(),
+      "dispatch_scaled_masked_softmax_backward",
+      dispatch_scaled_masked_softmax_backward<scalar_t, scalar_t, float>(
+          reinterpret_cast<scalar_t*>(output_grads_ptr), 
+	  reinterpret_cast<scalar_t*>(output_grads_ptr), 
+	  reinterpret_cast<scalar_t const*>(softmax_results.data_ptr()),
+	  scale_factor,
+	  query_seq_len,
+	  key_seq_len,
+	  batches,
+	  attn_heads);
+			   );
+  
+  //backward pass is completely in-place
+  return output_grads;
+}
+}
+}
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax.cpp
new file mode 100644
index 0000000000..29754fc590
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax.cpp
@@ -0,0 +1,72 @@
+/* coding=utf-8
+ * Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#include <cuda_fp16.h>
+#include <torch/extension.h>
+#include <vector>
+
+namespace multihead_attn {
+namespace fused_softmax {
+namespace scaled_upper_triang_masked_softmax {
+
+torch::Tensor fwd_cuda(
+    torch::Tensor const& input, 
+    float scale_factor);
+
+torch::Tensor bwd_cuda(
+    torch::Tensor const& output_grads, 
+    torch::Tensor const& softmax_results,
+    float scale_factor);
+
+torch::Tensor fwd(torch::Tensor const& input, float scale_factor) {
+  AT_ASSERTM(input.dim() == 3, "expected 3D tensor");
+  AT_ASSERTM((input.scalar_type() == at::ScalarType::Half) ||
+	     (input.scalar_type() == at::ScalarType::BFloat16), 
+      "Only fp16 and bf16 are supported");
+
+  return fwd_cuda(input, scale_factor);
+}
+
+torch::Tensor bwd(
+    torch::Tensor const& output_grads, 
+    torch::Tensor const& softmax_results,
+    float scale_factor) {
+
+  AT_ASSERTM(output_grads.dim() == 3, "expected 3D tensor");
+  AT_ASSERTM(softmax_results.dim() == 3, "expected 3D tensor");
+
+  AT_ASSERTM((output_grads.scalar_type() == at::ScalarType::Half) ||
+	     (output_grads.scalar_type() == at::ScalarType::BFloat16), 
+      "Only fp16 and bf16 are supported");
+  AT_ASSERTM((softmax_results.scalar_type() == at::ScalarType::Half) ||
+	     (softmax_results.scalar_type() == at::ScalarType::BFloat16), 
+      "Only fp16 and bf16 are supported");
+
+  return bwd_cuda(output_grads, softmax_results, scale_factor);
+}
+
+} // end namespace scaled_upper_triang_masked_softmax
+} // end namespace fused_softmax
+} // end namespace multihead_attn
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", 
+        &multihead_attn::fused_softmax::scaled_upper_triang_masked_softmax::fwd,
+	"Self Multihead Attention scaled, time masked softmax -- Forward.");
+  m.def("backward", 
+        &multihead_attn::fused_softmax::scaled_upper_triang_masked_softmax::bwd,
+	"Self Multihead Attention scaled, time masked softmax -- Backward.");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax.h
new file mode 100644
index 0000000000..445e0d88ce
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax.h
@@ -0,0 +1,513 @@
+/* coding=utf-8
+ * Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#pragma once
+
+#include <assert.h>
+#include <cuda_fp16.h>
+#include <cfloat>
+#include <limits>
+#include <stdint.h>
+#include <c10/macros/Macros.h>
+
+namespace {
+
+template <typename Datatype, int ELEMENTS_PER_LDG>
+__device__ __inline__ void copy_vector(Datatype *dst, const Datatype *src);
+
+template <>
+__device__ __inline__ void copy_vector<c10::BFloat16, 1>(c10::BFloat16 *dst, const c10::BFloat16 *src) { *dst = *src; }
+
+template <>
+__device__ __inline__ void copy_vector<c10::BFloat16, 4>(c10::BFloat16 *dst, const c10::BFloat16 *src) { *((float2*) dst) = *((float2*) src); }
+  
+template <>
+__device__ __inline__ void copy_vector<c10::Half, 1>(c10::Half *dst, const c10::Half *src) { *dst = *src; }
+
+template <>
+__device__ __inline__ void copy_vector<c10::Half, 4>(c10::Half *dst, const c10::Half *src) { *((float2*) dst) = *((float2*) src); }
+
+template <>
+__device__ __inline__ void copy_vector<uint8_t, 1>(uint8_t *dst, const uint8_t *src) { *dst = *src; }
+
+template <>
+__device__ __inline__ void copy_vector<uint8_t, 4>(uint8_t *dst, const uint8_t *src) {*((half2*) dst) = *((half2*) src); }
+
+template <typename Datatype, int ELEMENTS_PER_LDG>
+__device__ __inline__ void copy_zero_vector(Datatype *dst);
+
+template <>
+__device__ __inline__ void copy_zero_vector<c10::BFloat16, 1>(c10::BFloat16 *dst) { *dst = 0.0; }
+
+template <>
+__device__ __inline__ void copy_zero_vector<c10::BFloat16, 4>(c10::BFloat16 *dst) { *((float2*) dst) = make_float2(0.0f, 0.0f); }
+
+template <>
+__device__ __inline__ void copy_zero_vector<c10::Half, 1>(c10::Half *dst) { *dst = 0.0; }
+
+template <>
+__device__ __inline__ void copy_zero_vector<c10::Half, 4>(c10::Half *dst) { *((float2*) dst) = make_float2(0.0f, 0.0f); }
+
+
+int log2_ceil(int value) {
+    int log2_value = 0;
+    while ((1 << log2_value) < value) ++log2_value;
+    return log2_value;
+}
+
+template<typename T>
+struct Add {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a + b;
+  }
+};
+
+template<typename T>
+struct Max {
+  __device__ __forceinline__ T operator()(T a, T b) const {
+    return a < b ? b : a;
+  }
+};
+
+template <typename T>
+__device__ __forceinline__ T WARP_SHFL_XOR_NATIVE(T value, int laneMask, int width = warpSize, unsigned int mask = 0xffffffff)
+{
+#if CUDA_VERSION >= 9000
+    return __shfl_xor_sync(mask, value, laneMask, width);
+#else
+    return __shfl_xor(value, laneMask, width);
+#endif
+}
+
+template <typename acc_t, int WARP_BATCH, int WARP_SIZE, template<typename> class ReduceOp>
+__device__ __forceinline__ void warp_reduce(acc_t* sum) {
+    ReduceOp<acc_t> r;
+    #pragma unroll
+    for (int offset = WARP_SIZE / 2; offset > 0; offset /= 2) {
+        #pragma unroll
+        for (int i = 0;  i < WARP_BATCH;  ++i) {
+            acc_t b = WARP_SHFL_XOR_NATIVE(sum[i], offset, WARP_SIZE);
+            sum[i] = r(sum[i], b);
+        }
+    }
+}
+
+/*
+ * Extended softmax (from native aten pytorch) with following additional features
+ * 1) input scaling
+ * 2) Implicit time (diagonal masking)
+ */
+template <typename input_t, typename output_t, typename acc_t, int log2_elements>
+__global__ void scaled_upper_triang_masked_softmax_warp_forward(
+    output_t *dst, 
+    const input_t *src, 
+    const acc_t scale, 
+    int micro_batch_size, 
+    int stride, 
+    int element_count) 
+{
+    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and 
+    // warp_size of method warp_softmax_forward_kernel.
+    constexpr int next_power_of_two = 1 << log2_elements;
+    constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+    constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
+    constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
+    constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;
+
+    int first_batch = (blockDim.y * blockIdx.y + threadIdx.y) * gridDim.x * WARP_BATCH + blockIdx.x;
+    int local_seq = blockIdx.x + 1; 
+    int warp_iteration_limit = (local_seq + ELEMENTS_PER_LDG_STG * WARP_SIZE - 1)/ WARP_SIZE;
+
+    // micro_batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = micro_batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+
+    src += first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+    dst += first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+
+    // load data from global memory
+    acc_t elements[WARP_BATCH][WARP_ITERATIONS];
+    input_t temp_data[ELEMENTS_PER_LDG_STG];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : local_seq;
+
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  it+=ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+
+            if (element_index < batch_element_count) {
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(temp_data, src + i*element_count*stride + it*WARP_SIZE);
+
+                #pragma unroll
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    if ((element_index + element) < batch_element_count) {
+                        elements[i][it+element] = (acc_t)temp_data[element] * scale;
+                    } else {
+                        elements[i][it + element] = -std::numeric_limits<acc_t>::infinity();
+                    }
+                }
+            } else {
+                #pragma unroll
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    elements[i][it + element] = -std::numeric_limits<acc_t>::infinity();
+                }
+            }
+        }
+    }
+
+    // compute max_value
+    acc_t max_value[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        max_value[i] = elements[i][0];
+        #pragma unroll
+        for (int it = 1;  it < WARP_ITERATIONS;  ++it) {
+            max_value[i] = (max_value[i] > elements[i][it]) ? max_value[i] : elements[i][it];
+        }
+    }
+    warp_reduce<acc_t, WARP_BATCH, WARP_SIZE, Max>(max_value);
+
+    acc_t sum[WARP_BATCH] { 0.0f };
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  ++it) {
+            if (it < warp_iteration_limit) {
+                elements[i][it] = std::exp((elements[i][it] - max_value[i]));
+                sum[i] += elements[i][it];
+            } 
+        }
+    }
+    warp_reduce<acc_t, WARP_BATCH, WARP_SIZE, Add>(sum);
+
+    // store result
+    output_t out[ELEMENTS_PER_LDG_STG];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  it+=ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+
+            if (element_index < local_seq) {
+
+                #pragma unroll  
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    if (element_index + element < local_seq) {
+                        out[element] = elements[i][it + element] / sum[i];
+                    } else {
+                        out[element] = 0;
+                    }
+                }
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(dst + i * element_count * stride + it * WARP_SIZE, out);
+            } else if (element_index < element_count) {
+                copy_zero_vector<output_t, ELEMENTS_PER_LDG_STG>(dst + i * element_count * stride + it * WARP_SIZE);
+            } else {
+                break;
+            } 
+        }
+    }
+}
+
+template <typename input_t, typename output_t, typename acc_t, int log2_elements>
+__global__ void scaled_upper_triang_masked_softmax_warp_backward(
+    output_t *gradInput, 
+    input_t *grad, 
+    const input_t *output,
+    acc_t scale, 
+    int micro_batch_size, 
+    int stride, 
+    int element_count)
+{
+    // WARP_SIZE and WARP_BATCH must match the return values batches_per_warp and 
+    // warp_size of method warp_softmax_backward_kernel.
+    constexpr int next_power_of_two = 1 << log2_elements;
+    constexpr int WARP_SIZE = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+    constexpr int WARP_ITERATIONS = next_power_of_two / WARP_SIZE;
+    constexpr int WARP_BATCH = (next_power_of_two <= 128) ? 2 : 1;
+    constexpr int ELEMENTS_PER_LDG_STG = (WARP_ITERATIONS < 4) ? 1 : 4;
+
+    int first_batch = (blockDim.y * blockIdx.y + threadIdx.y) * gridDim.x * WARP_BATCH + blockIdx.x;
+    int local_seq = blockIdx.x + 1; 
+    
+    // micro_batch_size might not be a multiple of WARP_BATCH. Check how
+    // many batches have to computed within this WARP.
+    int local_batches = micro_batch_size - first_batch;
+    if (local_batches > WARP_BATCH)
+        local_batches = WARP_BATCH;
+
+    // there might be multiple batches per warp. compute the index within the batch
+    int local_idx = threadIdx.x;
+
+    // the first element to process by the current thread
+    int thread_offset = first_batch * stride + ELEMENTS_PER_LDG_STG * local_idx;
+    grad += thread_offset;
+    output += thread_offset;
+    gradInput += thread_offset;
+
+    // load data from global memory
+    acc_t grad_reg[WARP_BATCH][WARP_ITERATIONS] { 0.0f };
+    acc_t output_reg[WARP_BATCH][WARP_ITERATIONS] { 0.0f };
+    input_t temp_grad[ELEMENTS_PER_LDG_STG];
+    input_t temp_output[ELEMENTS_PER_LDG_STG];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        int batch_element_count = (i >= local_batches) ? 0 : local_seq;
+
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  it+=ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < batch_element_count) {
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(temp_grad, grad + i * element_count * stride + it * WARP_SIZE);
+                copy_vector<input_t, ELEMENTS_PER_LDG_STG>(temp_output, output + i * element_count * stride + it * WARP_SIZE);
+
+                #pragma unroll
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    if (element_index + element < batch_element_count) {
+                        output_reg[i][it + element] = (acc_t)temp_output[element];
+                    }
+                }
+                #pragma unroll
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    if (element_index + element < batch_element_count) {
+                        grad_reg[i][it + element] = (acc_t)temp_grad[element] * output_reg[i][it + element];
+                    }
+                }
+            }
+        }
+    }
+   
+    acc_t sum[WARP_BATCH];
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        sum[i] = grad_reg[i][0];
+        #pragma unroll
+        for (int it = 1;  it < WARP_ITERATIONS;  ++it) {
+            sum[i] += grad_reg[i][it];
+        }
+    }
+    warp_reduce<acc_t, WARP_BATCH, WARP_SIZE, Add>(sum);
+
+    // store result
+    #pragma unroll
+    for (int i = 0;  i < WARP_BATCH;  ++i) {
+        if (i >= local_batches)
+            break;
+        #pragma unroll
+        for (int it = 0;  it < WARP_ITERATIONS;  it+=ELEMENTS_PER_LDG_STG) {
+            int element_index = ELEMENTS_PER_LDG_STG * local_idx + it * WARP_SIZE;
+            if (element_index < element_count) {
+                // compute gradients
+                output_t out[ELEMENTS_PER_LDG_STG];
+                #pragma unroll
+                for (int element = 0; element < ELEMENTS_PER_LDG_STG; ++element) {
+                    out[element] = (output_t)(scale * (grad_reg[i][it + element] - output_reg[i][it + element] * sum[i]));
+                }
+                copy_vector<output_t, ELEMENTS_PER_LDG_STG>(gradInput + i * element_count * stride + it * WARP_SIZE, out);
+            } 
+        }
+    }
+}
+
+} // end of anonymous namespace
+
+template<typename input_t, typename output_t, typename acc_t>
+void dispatch_scaled_upper_triang_masked_softmax_forward(
+    output_t *dst, 
+    const input_t *src, 
+    const input_t scale, 
+    int softmax_elements, 
+    int softmax_elements_stride, 
+    int attn_batches)
+{
+    TORCH_INTERNAL_ASSERT(softmax_elements >= 0 && softmax_elements <= 2048 );
+    if (softmax_elements == 0) {
+        return;
+    } else {
+        int log2_elements = log2_ceil(softmax_elements);
+        const int next_power_of_two = 1 << log2_elements;
+        int seq_len = softmax_elements;
+        int batch_count = attn_batches * seq_len;
+
+        // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_forward.
+        int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+
+        // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_forward.
+        int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+
+        int warps_per_block = (threads_per_block / warp_size);
+        int batches_per_block = warps_per_block * batches_per_warp;
+        TORCH_INTERNAL_ASSERT(attn_batches % batches_per_block == 0);
+
+        int blocks_per_seq = attn_batches / batches_per_block;
+        dim3 blocks(seq_len, blocks_per_seq, 1);
+        dim3 threads(warp_size, warps_per_block, 1);
+        // Launch code would be more elegant if C++ supported FOR CONSTEXPR
+        switch (log2_elements) {
+            case 0: // 1
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 0>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 1: // 2
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 1>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 2: // 4
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 2>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 3: // 8
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 3>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 4: // 16
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 4>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 5: // 32
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 5>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 6: // 64
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 6>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 7: // 128
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 7>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 8: // 256
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 8>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 9: // 512
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 9>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 10: // 1024
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 10>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 11: // 2048
+                scaled_upper_triang_masked_softmax_warp_forward<input_t, output_t, acc_t, 11>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(dst, src, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            default:
+                break;
+        }
+    }
+}
+
+template<typename input_t, typename output_t, typename acc_t>
+void dispatch_scaled_upper_triang_masked_softmax_backward(
+    output_t *grad_input, 
+    input_t *grad, 
+    const input_t *output, 
+    const acc_t scale, 
+    int softmax_elements, 
+    int softmax_elements_stride, 
+    int attn_batches)
+{
+    TORCH_INTERNAL_ASSERT( softmax_elements >= 0 && softmax_elements <= 2048 );
+    if (softmax_elements == 0) {
+       return;
+    } else {
+        int log2_elements = log2_ceil(softmax_elements);
+        const int next_power_of_two = 1 << log2_elements;
+        int seq_len = softmax_elements;
+        int batch_count = attn_batches * seq_len;
+
+        // This value must match the WARP_SIZE constexpr value computed inside softmax_warp_backward.
+        int warp_size = (next_power_of_two < C10_WARP_SIZE) ? next_power_of_two : C10_WARP_SIZE;
+
+        // This value must match the WARP_BATCH constexpr value computed inside softmax_warp_backward.
+        int batches_per_warp = (next_power_of_two <= 128) ? 2 : 1;
+
+        // use 128 threads per block to maximimize gpu utilization
+        constexpr int threads_per_block = 128;
+
+        int warps_per_block = (threads_per_block / warp_size);
+        int batches_per_block = warps_per_block * batches_per_warp;
+        TORCH_INTERNAL_ASSERT(attn_batches % batches_per_block == 0);
+
+        int blocks_per_seq = attn_batches / batches_per_block;
+        dim3 blocks(seq_len, blocks_per_seq, 1);
+        dim3 threads(warp_size, warps_per_block, 1);
+        // Launch code would be more elegant if C++ supported FOR CONSTEXPR
+        switch (log2_elements) {
+            case 0: // 1
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 0>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 1: // 2
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 1>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 2: // 4
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 2>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 3: // 8
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 3>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 4: // 16
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 4>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 5: // 32
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 5>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 6: // 64
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 6>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 7: // 128
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 7>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 8: // 256
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 8>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 9: // 512
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 9>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 10: // 1024
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 10>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            case 11: // 2048
+                scaled_upper_triang_masked_softmax_warp_backward<input_t, output_t, acc_t, 11>
+                    <<<blocks, threads, 0, at::cuda::getCurrentCUDAStream()>>>(grad_input, grad, output, scale, batch_count, softmax_elements_stride, softmax_elements);
+                break;
+            default:
+                break;
+        }
+    }
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu
new file mode 100644
index 0000000000..a90a9344f3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu
@@ -0,0 +1,98 @@
+/* coding=utf-8
+ * Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+ *
+ * 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.
+ */
+
+#include <ATen/ATen.h>
+#include <cuda.h>
+#include <cuda_runtime.h>
+#include <cuda_fp16.h>
+#include <cuda_profiler_api.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <torch/extension.h>
+#include "scaled_upper_triang_masked_softmax.h"
+#include "type_shim.h"
+
+namespace multihead_attn {
+namespace fused_softmax {
+namespace scaled_upper_triang_masked_softmax {
+
+torch::Tensor fwd_cuda(
+    torch::Tensor const& input, 
+    float scale_factor)
+{
+  // input is a 3d tensor with dimensions [attn_batches, seq_len, seq_len]
+  const int attn_batches = input.size(0);
+  const int seq_len = input.size(1);
+  TORCH_INTERNAL_ASSERT(seq_len <= 2048);
+
+  // Output 
+  auto act_options = input.options().requires_grad(false);
+  torch::Tensor softmax_results = 
+      torch::empty({attn_batches, seq_len, seq_len}, act_options);
+
+  // Softmax Intermediate Result Ptr
+  void* input_ptr = static_cast<void*>(input.data_ptr());
+  void* softmax_results_ptr = static_cast<void*>(softmax_results.data_ptr());
+
+  DISPATCH_HALF_AND_BFLOAT(
+      input.scalar_type(),
+      "dispatch_scaled_upper_triang_masked_softmax_forward",
+      dispatch_scaled_upper_triang_masked_softmax_forward<scalar_t, scalar_t, float>(
+	  reinterpret_cast<scalar_t*>(softmax_results_ptr),
+	  reinterpret_cast<const scalar_t*>(input_ptr),
+	  scale_factor,
+	  seq_len,
+	  seq_len,
+	  attn_batches);
+      );
+  return softmax_results;
+}
+				      
+
+torch::Tensor bwd_cuda(
+    torch::Tensor const& output_grads_, 
+    torch::Tensor const& softmax_results_, 
+    float scale_factor)  {
+	
+  auto output_grads = output_grads_.contiguous();
+  auto softmax_results = softmax_results_.contiguous();
+
+  //output grads is a 3d tensor with dimensions [attn_batches, seq_len, seq_len]
+  const int attn_batches = output_grads.size(0);
+  const int seq_len = output_grads.size(1);
+  TORCH_INTERNAL_ASSERT(output_grads.size(1) == output_grads.size(2));
+
+  void* output_grads_ptr = static_cast<void*>(output_grads.data_ptr());
+
+  //Softmax Grad
+  DISPATCH_HALF_AND_BFLOAT(
+      output_grads_.scalar_type(),
+      "dispatch_scaled_upper_triang_masked_softmax_backward",
+      dispatch_scaled_upper_triang_masked_softmax_backward<scalar_t, scalar_t, float>(
+          reinterpret_cast<scalar_t*>(output_grads_ptr), 
+	  reinterpret_cast<scalar_t*>(output_grads_ptr), 
+	  reinterpret_cast<scalar_t const*>(softmax_results.data_ptr()),
+	  scale_factor,
+	  seq_len,
+	  seq_len,
+	  attn_batches);
+      );
+  
+  //backward pass is completely in-place
+  return output_grads;
+}
+}
+}
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/mlp.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/mlp.cpp
new file mode 100644
index 0000000000..68cf174c22
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/mlp.cpp
@@ -0,0 +1,166 @@
+#include <torch/extension.h>
+#include <torch/torch.h>
+#include <vector>
+
+#include <stdio.h>
+
+size_t get_mlp_reserved_space(int64_t batch_size, int num_layers, const int* output_features);
+
+template <typename T>
+size_t get_mlp_bp_workspace_in_bytes(int batch_size, int num_layers, const int* output_features);
+
+template <typename T>
+int mlp_fp(
+    T* X,
+    int input_features,
+    int batch_size,
+    T** WPtr,
+    int num_layers,
+    int* output_features,
+    T** BPtr,
+    T* Y,
+    T* reserved_space,
+    int use_bias,
+    int activation,
+    void* lt_workspace);
+
+template <typename T>
+int mlp_bp(
+    T* X,
+    T* Y,
+    int input_features,
+    int batch_size,
+    T** WPtr,
+    int num_layers,
+    int* output_features,
+    T* dY,
+    T* reserved_space,
+    T* work_space,
+    T* dX,
+    T** dwPtr,
+    T** dbPtr,
+    bool requires_grad,
+    int use_bias,
+    int activation);
+
+std::vector<at::Tensor> mlp_forward(int use_bias, int activation, std::vector<at::Tensor> inputs) {
+
+  auto num_layers = inputs.size() - 1;
+  if (use_bias) {
+    // inputs contains (input, weights, biases)
+    num_layers /= 2;
+  }
+  auto batch_size = inputs[0].size(0);
+  auto input_features = inputs[0].size(1);
+
+  std::vector<int> output_features;
+  for (int i = 0; i < num_layers; i++) {
+    output_features.push_back(inputs[i + 1].size(0));
+  }
+
+  auto reserved_size = get_mlp_reserved_space(batch_size, num_layers, output_features.data());
+
+  // create output/workspace tensor
+  auto out = at::empty({batch_size, output_features.back()}, inputs[0].type());
+  auto reserved_space = at::empty({reserved_size}, inputs[0].type());
+  // allocate fixed 4MB workspace for cublaslt for now, and this gets at least 4 MB
+  auto lt_workspace = at::empty({1 << 22}, inputs[0].type());
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(inputs[0].type(), "mlp_forward", [&] {
+    std::vector<scalar_t*> w_ptr;
+    std::vector<scalar_t*> b_ptr;
+    for (int i = 0; i < num_layers; i++) {
+      w_ptr.push_back(inputs[i + 1].data_ptr<scalar_t>());
+      if (use_bias) {
+        b_ptr.push_back(inputs[i + 1 + num_layers].data_ptr<scalar_t>());
+      }
+    }
+    auto result = mlp_fp<scalar_t>(
+        inputs[0].data_ptr<scalar_t>(),
+        input_features,
+        batch_size,
+        w_ptr.data(),
+        num_layers,
+        output_features.data(),
+        b_ptr.data(),
+        out.data_ptr<scalar_t>(),
+        reserved_space.data_ptr<scalar_t>(),
+        use_bias,
+        activation,
+        (void*) (lt_workspace.data_ptr<scalar_t>()));
+  });
+
+  return {out, reserved_space};
+}
+
+std::vector<at::Tensor> mlp_backward(
+  int use_bias,
+  int activation,
+  at::Tensor grad_o,
+  std::vector<at::Tensor> fprop_outputs,
+  std::vector<at::Tensor> inputs) {
+
+  auto num_layers = inputs.size() - 1;
+  if (use_bias) {
+    // inputs contains (input, weights, biases)
+    num_layers /= 2;
+  }
+
+  auto batch_size = inputs[0].size(0);
+  auto input_features = inputs[0].size(1);
+
+  bool requires_grad = inputs[0].requires_grad();
+
+  std::vector<int> output_features;
+  for (int i = 0; i < num_layers; i++) {
+    output_features.push_back(inputs[i + 1].size(0));
+  }
+  // create outputs, length of inputs
+  std::vector<at::Tensor> outputs;
+  for (int i = 0; i < inputs.size(); i++) {
+    outputs.push_back(at::empty(inputs[i].sizes(), inputs[i].type()));  // clone for testing now
+  }
+
+  AT_DISPATCH_FLOATING_TYPES_AND_HALF(inputs[0].type(), "mlp_backward", [&] {
+    std::vector<scalar_t*> w_ptr;
+    for (int i = 0; i < num_layers; i++) {
+      w_ptr.push_back(inputs[i + 1].data_ptr<scalar_t>());
+    }
+    std::vector<scalar_t*> outputs_ptr;
+    for (int i = 0; i < inputs.size(); i++) {
+      outputs_ptr.push_back(outputs[i].data_ptr<scalar_t>());
+    }
+
+    auto work_size =
+        get_mlp_bp_workspace_in_bytes<scalar_t>(batch_size, num_layers, output_features.data());
+
+    // auto work_space = at::empty({work_size*4}, at::kByte);
+    auto work_space = at::empty({work_size / sizeof(scalar_t)}, inputs[0].type());
+
+    auto result = mlp_bp<scalar_t>(
+        inputs[0].data_ptr<scalar_t>(),
+        fprop_outputs[0].data_ptr<scalar_t>(),
+        input_features,
+        batch_size,
+        w_ptr.data(),
+        num_layers,
+        output_features.data(),
+        grad_o.contiguous().data_ptr<scalar_t>(),
+        fprop_outputs[1].data_ptr<scalar_t>(),
+        work_space.data_ptr<scalar_t>(),
+        outputs_ptr[0],
+        outputs_ptr.data() + 1,
+        outputs_ptr.data() + 1 + num_layers,
+        requires_grad,
+        use_bias,
+        activation);
+  });
+
+  return outputs;
+}
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("forward", &mlp_forward, "MLP forward");
+  m.def("backward", &mlp_backward, "MLP backward");
+}
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/mlp_cuda.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/mlp_cuda.cu
new file mode 100644
index 0000000000..f93f1df1a2
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/mlp_cuda.cu
@@ -0,0 +1,1678 @@
+#include <ATen/ATen.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <assert.h>
+#include <stdio.h>
+#include <stdlib.h>
+#include <string.h>
+#include <torch/torch.h>
+
+/* Includes, cuda */
+#include <cublas_v2.h>
+#include <cuda_runtime.h>
+
+#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000
+// includes cublaslt
+#include <cublasLt.h>
+#endif
+// constants for fused bias+relu kernel
+#define BIAS_RELU_FW_NTHREADS 128 // forward number of thread per block
+#define BIAS_RELU_BW_NTHREADS_X 32 // backward number of thread in feature dim
+#define BIAS_RELU_BW_NTHREADS_Y 16 // backward number of thread in batch dim
+#define BIAS_RELU_RED_PER_THREAD 16 // backward minimal reduction length per thread
+
+// move to a header later on
+#define ILP 4
+template<typename T>
+__host__ __device__ __forceinline__ bool is_aligned(T* p){
+  return ((uint64_t)p) % (ILP*sizeof(T)) == 0;
+}
+
+template<typename T>
+__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+template<typename T>
+__device__ __forceinline__ void load_store(T* dst, volatile T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+template<typename T>
+__device__ __forceinline__ void load_store(volatile T* dst, T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+
+// Keep ReLU in float only. When using half, cast to float before calling.
+__device__ __inline__ float relu(float a) {
+  float retf = max(a, 0.f);
+  return (retf);
+}
+
+// Keep Sigmoid in float only. When using half, cast to float before calling.
+__device__ __inline__ float sigmoid(float a) {
+  float retf = 1.f / (1.f + expf(-a));
+  return (retf);
+}
+
+// FP64 Wrapper around cublas GEMMEx
+cublasStatus_t mlp_gemm(
+    cublasHandle_t handle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    float* alpha,
+    const double* A,
+    int lda,
+    const double* B,
+    int ldb,
+    const float* beta,
+    double* C,
+    int ldc) {
+  return cublasGemmEx(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      CUDA_R_64F,
+      lda,
+      B,
+      CUDA_R_64F,
+      ldb,
+      beta,
+      C,
+      CUDA_R_64F,
+      ldc,
+      CUDA_R_64F,
+      CUBLAS_GEMM_DEFAULT);
+}
+
+// FP32 Wrapper around cublas GEMMEx
+cublasStatus_t mlp_gemm(
+    cublasHandle_t handle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    float* alpha,
+    const float* A,
+    int lda,
+    const float* B,
+    int ldb,
+    const float* beta,
+    float* C,
+    int ldc) {
+  return cublasGemmEx(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      CUDA_R_32F,
+      lda,
+      B,
+      CUDA_R_32F,
+      ldb,
+      beta,
+      C,
+      CUDA_R_32F,
+      ldc,
+      CUDA_R_32F,
+      CUBLAS_GEMM_DEFAULT);
+}
+
+// FP16 Tensor core wrapper around cublas GEMMEx
+cublasStatus_t mlp_gemm(
+    cublasHandle_t handle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    float* alpha,
+    const at::Half* A,
+    int lda,
+    const at::Half* B,
+    int ldb,
+    float* beta,
+    at::Half* C,
+    int ldc) {
+  return cublasGemmEx(
+      handle,
+      transa,
+      transb,
+      m,
+      n,
+      k,
+      alpha,
+      A,
+      CUDA_R_16F,
+      lda,
+      B,
+      CUDA_R_16F,
+      ldb,
+      beta,
+      C,
+      CUDA_R_16F,
+      ldc,
+      CUDA_R_32F,
+      CUBLAS_GEMM_DEFAULT_TENSOR_OP);
+}
+#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000
+int mlp_gemm_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    float *alpha, /* host pointer */
+    const at::Half* A,
+    int lda,
+    const at::Half* B,
+    int ldb,
+    float *beta, /* host pointer */
+    at::Half* C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    bool use_relu,
+    const void* bias) {
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  cublasLtMatmulDescOpaque_t operationDesc = {};
+  cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};
+  cublasLtMatmulPreferenceOpaque_t preference = {};
+
+  int returnedResults                             = 0;
+  cublasLtMatmulHeuristicResult_t heuristicResult = {};
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;
+
+  // Create operation descriptor; see cublasLtMatmulDescAttributes_t
+  // for details about defaults; here we just set the transforms for
+  // A and B.
+  status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (use_bias) {
+    status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias));
+    if (status != CUBLAS_STATUS_SUCCESS) {
+      goto CLEANUP;
+    }
+    if (use_relu) {
+      epilogue = CUBLASLT_EPILOGUE_RELU_BIAS;
+    } else {
+      epilogue = CUBLASLT_EPILOGUE_BIAS;
+    }
+  } else {
+    if (use_relu) {
+      epilogue = CUBLASLT_EPILOGUE_RELU;
+    }
+  }
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+
+  // Create matrix descriptors. Not setting any extra attributes.
+  status = cublasLtMatrixLayoutInit(
+    &Adesc, CUDA_R_16F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(
+    &Bdesc, CUDA_R_16F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_16F, m, n, ldc);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // Create preference handle; In general, extra attributes can be
+  // used here to disable tensor ops or to make sure algo selected
+  // will work with badly aligned A, B, C. However, for simplicity
+  // here we assume A,B,C are always well aligned (e.g., directly
+  // come from cudaMalloc)
+  status = cublasLtMatmulPreferenceInit(&preference);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulPreferenceSetAttribute(
+    &preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // We just need the best available heuristic to try and run matmul.
+  // There is no guarantee that this will work. For example, if A is
+  // badly aligned, you can request more (e.g. 32) algos and try to
+  // run them one by one until something works.
+  status = cublasLtMatmulAlgoGetHeuristic(
+    ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1, &heuristicResult, &returnedResults);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (returnedResults == 0) {
+    status = CUBLAS_STATUS_NOT_SUPPORTED;
+    goto CLEANUP;
+  }
+  status = cublasLtMatmul(ltHandle,
+                          &operationDesc,
+                          alpha,
+                          A,
+                          &Adesc,
+                          B,
+                          &Bdesc,
+                          beta,
+                          C,
+                          &Cdesc,
+                          C,
+                          &Cdesc,
+                          &heuristicResult.algo,
+                          workspace,
+                          workspaceSize,
+                          stream);
+
+CLEANUP:
+  // Descriptors are no longer needed as all GPU work was already
+  // enqueued.
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;
+}
+
+int mlp_gemm_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    float *alpha, /* host pointer */
+    const double* A,
+    int lda,
+    const double* B,
+    int ldb,
+    float *beta, /* host pointer */
+    double* C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    bool use_relu,
+    const void* bias) {
+  return 1;
+}
+
+int mlp_gemm_lt(
+    cublasLtHandle_t ltHandle,
+    cublasOperation_t transa,
+    cublasOperation_t transb,
+    int m,
+    int n,
+    int k,
+    float *alpha, /* host pointer */
+    const float *A,
+    int lda,
+    const float *B,
+    int ldb,
+    float *beta, /* host pointer */
+    float *C,
+    int ldc,
+    void *workspace,
+    size_t workspaceSize,
+    cudaStream_t stream,
+    bool use_bias,
+    bool use_relu,
+    const void* bias) {
+  cublasStatus_t status = CUBLAS_STATUS_SUCCESS;
+
+  cublasLtMatmulDescOpaque_t operationDesc = {};
+  cublasLtMatrixLayoutOpaque_t Adesc = {}, Bdesc = {}, Cdesc = {};
+  cublasLtMatmulPreferenceOpaque_t preference = {};
+
+  int returnedResults                             = 0;
+  cublasLtMatmulHeuristicResult_t heuristicResult = {};
+  cublasLtEpilogue_t epilogue = CUBLASLT_EPILOGUE_DEFAULT;
+
+  // Create operation descriptor; see cublasLtMatmulDescAttributes_t
+  // for details about defaults; here we just set the transforms for
+  // A and B.
+  status = cublasLtMatmulDescInit(&operationDesc, CUBLAS_COMPUTE_32F, CUDA_R_32F);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSA, &transa, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_TRANSB, &transb, sizeof(transa));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (use_bias) {
+    status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_BIAS_POINTER, &bias, sizeof(bias));
+    if (status != CUBLAS_STATUS_SUCCESS) {
+      goto CLEANUP;
+    }
+    if (use_relu) {
+      epilogue = CUBLASLT_EPILOGUE_RELU_BIAS;
+    } else {
+      epilogue = CUBLASLT_EPILOGUE_BIAS;
+    }
+  } else {
+    if (use_relu) {
+      epilogue = CUBLASLT_EPILOGUE_RELU;
+    }
+  }
+
+  status = cublasLtMatmulDescSetAttribute(&operationDesc, CUBLASLT_MATMUL_DESC_EPILOGUE, &epilogue, sizeof(epilogue));
+  if (status != CUBLAS_STATUS_SUCCESS) {
+    goto CLEANUP;
+  }
+
+  // Create matrix descriptors. Not setting any extra attributes.
+  status = cublasLtMatrixLayoutInit(
+    &Adesc, CUDA_R_32F, transa == CUBLAS_OP_N ? m : k, transa == CUBLAS_OP_N ? k : m, lda);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(
+    &Bdesc, CUDA_R_32F, transb == CUBLAS_OP_N ? k : n, transb == CUBLAS_OP_N ? n : k, ldb);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatrixLayoutInit(&Cdesc, CUDA_R_32F, m, n, ldc);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // Create preference handle; In general, extra attributes can be
+  // used here to disable tensor ops or to make sure algo selected
+  // will work with badly aligned A, B, C. However, for simplicity
+  // here we assume A,B,C are always well aligned (e.g., directly
+  // come from cudaMalloc)
+  status = cublasLtMatmulPreferenceInit(&preference);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+  status = cublasLtMatmulPreferenceSetAttribute(
+    &preference, CUBLASLT_MATMUL_PREF_MAX_WORKSPACE_BYTES, &workspaceSize, sizeof(workspaceSize));
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  // We just need the best available heuristic to try and run matmul.
+  // There is no guarantee that this will work. For example, if A is
+  // badly aligned, you can request more (e.g. 32) algos and try to
+  // run them one by one until something works.
+  status = cublasLtMatmulAlgoGetHeuristic(
+    ltHandle, &operationDesc, &Adesc, &Bdesc, &Cdesc, &Cdesc, &preference, 1, &heuristicResult, &returnedResults);
+  if (status != CUBLAS_STATUS_SUCCESS) goto CLEANUP;
+
+  if (returnedResults == 0) {
+    status = CUBLAS_STATUS_NOT_SUPPORTED;
+    goto CLEANUP;
+  }
+
+  status = cublasLtMatmul(ltHandle,
+                          &operationDesc,
+                          alpha,
+                          A,
+                          &Adesc,
+                          B,
+                          &Bdesc,
+                          beta,
+                          C,
+                          &Cdesc,
+                          C,
+                          &Cdesc,
+                          &heuristicResult.algo,
+                          workspace,
+                          workspaceSize,
+                          stream);
+
+CLEANUP:
+  // Descriptors are no longer needed as all GPU work was already
+  // enqueued.
+  return status == CUBLAS_STATUS_SUCCESS ? 0 : 1;
+}
+#endif
+
+// Bias ADD. Assume input X is [features x batch size], column major.
+// Bias is one 'features' long vector, with implicit broadcast.
+template <typename T>
+__global__ void biasAdd_fprop(T *X, T *b, uint batch_size, uint features) {
+  T r_x[ILP];
+  T r_b[ILP];
+  if(is_aligned(X) && is_aligned(b) && features % ILP ==0) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid*ILP < features * batch_size; tid += blockDim.x * gridDim.x) {
+      int row = tid % (features / ILP);
+      load_store(r_x, X, 0 , tid);
+      load_store(r_b, b, 0 , row);
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        float bias_sum = static_cast<float>(r_x[ii]) + static_cast<float>(r_b[ii]);
+        r_x[ii] = bias_sum;
+      }
+      load_store(X, r_x, tid , 0);
+    }
+  } else {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          int row = tid % features;
+          r_x[ii] = X[idx];
+          r_b[ii] = b[row];
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        float bias_sum = static_cast<float>(r_x[ii]) + static_cast<float>(r_b[ii]);
+        r_x[ii] = bias_sum;
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          X[idx] = r_x[ii];
+        }
+      }
+    }
+  }
+}
+
+// Bias ADD + ReLU. Assume input X is [features x batch size], column major.
+// Activation support fuesed ReLU. Safe to call in-place.
+template <typename T>
+__global__ void biasAddRelu_fprop(T *X, T *b, uint batch_size, uint features) {
+  T r_x[ILP];
+  T r_b[ILP];
+  if(is_aligned(X) && is_aligned(b) && features % ILP ==0) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid*ILP < features * batch_size; tid += blockDim.x * gridDim.x) {
+      int row = tid % (features / ILP);
+      load_store(r_x, X, 0 , tid);
+      load_store(r_b, b, 0 , row);
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        float bias_sum = static_cast<float>(r_x[ii]) + static_cast<float>(r_b[ii]);
+        r_x[ii] = relu(bias_sum);
+      }
+      load_store(X, r_x, tid , 0);
+    }
+  } else {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          int row = tid % features;
+          r_x[ii] = X[idx];
+          r_b[ii] = b[row];
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        float bias_sum = static_cast<float>(r_x[ii]) + static_cast<float>(r_b[ii]);
+        r_x[ii] = relu(bias_sum);
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          X[idx] = r_x[ii];
+        }
+      }
+    }
+  }
+}
+
+// ReLU. Assume input X is [features x batch size], column major.
+// Safe to call in-place.
+template <typename T>
+__global__ void Relu_fprop(T *X, uint batch_size, uint features) {
+  T r_x[ILP];
+  if(is_aligned(X) && features % ILP ==0) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid*ILP < features * batch_size; tid += blockDim.x * gridDim.x) {
+      load_store(r_x, X, 0 , tid);
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        r_x[ii] = relu(static_cast<float>(r_x[ii]));
+      }
+      load_store(X, r_x, tid , 0);
+    }
+  } else {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          r_x[ii] = X[idx];
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        r_x[ii] = relu(static_cast<float>(r_x[ii]));
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          X[idx] = r_x[ii];
+        }
+      }
+    }
+  }
+}
+
+// Sigmoid. Assume input X is [features x batch size], column major.
+// Safe to call in-place.
+template <typename T>
+__global__ void Sigmoid_fprop(T *X, uint batch_size, uint features) {
+  T r_x[ILP];
+  if(is_aligned(X) && features % ILP ==0) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid*ILP < features * batch_size; tid += blockDim.x * gridDim.x) {
+      load_store(r_x, X, 0 , tid);
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        r_x[ii] = sigmoid(static_cast<float>(r_x[ii]));
+      }
+      load_store(X, r_x, tid , 0);
+    }
+  } else {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          r_x[ii] = X[idx];
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        r_x[ii] = sigmoid(static_cast<float>(r_x[ii]));
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          X[idx] = r_x[ii];
+        }
+      }
+    }
+  }
+}
+
+// ReLU. Assume input X is [features x batch size], column major.
+// Safe to call in-place.
+template <typename T>
+__global__ void Relu_bprop(T *dY, T *Y, uint batch_size, uint features, T *dX) {
+  T r_dy[ILP];
+  T r_y[ILP];
+  if(is_aligned(dY) &&
+     is_aligned(Y) &&
+     is_aligned(dX) &&
+     features % ILP ==0) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid*ILP < features * batch_size; tid += blockDim.x * gridDim.x) {
+      load_store(r_dy, dY, 0 , tid);
+      load_store(r_y, Y, 0 , tid);
+#pragma unroll
+      for(int ii=0;ii<ILP;ii++){
+        if ((float)r_y[ii] <= 0.f)
+          r_dy[ii] = 0;
+      }
+      load_store(dX, r_dy, tid, 0);
+    }
+  } else {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          r_dy[ii] = dY[idx];
+          r_y[ii] = Y[idx];
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        if ((float)r_y[ii] <= 0.f)
+          r_dy[ii] = 0;
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          dX[idx] = r_dy[ii];
+        }
+      }
+    }
+  }
+}
+
+// Sigmoid. Assume input X is [features x batch size], column major.
+// Safe to call in-place.
+template <typename T>
+__global__ void Sigmoid_bprop(T *dY, T *Y, uint batch_size, uint features, T *dX) {
+  T r_dy[ILP];
+  T r_y[ILP];
+  if(is_aligned(dY) &&
+     is_aligned(Y) &&
+     is_aligned(dX) &&
+     features % ILP ==0) {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid*ILP < features * batch_size; tid += blockDim.x * gridDim.x) {
+      load_store(r_dy, dY, 0 , tid);
+      load_store(r_y, Y, 0 , tid);
+#pragma unroll
+      for(int ii=0;ii<ILP;ii++){
+        float grad_out = r_dy[ii];
+        float out = r_y[ii];
+        float grad_i = out * ( 1.f - out) * grad_out;
+        r_dy[ii] = grad_i;
+      }
+      load_store(dX, r_dy, tid, 0);
+    }
+  } else {
+    int tid = blockIdx.x * blockDim.x + threadIdx.x;
+    for (; tid < features * batch_size; tid += ILP * blockDim.x * gridDim.x) {
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          r_dy[ii] = dY[idx];
+          r_y[ii] = Y[idx];
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        float grad_out = r_dy[ii];
+        float out = r_y[ii];
+        float grad_i = out * ( 1.f - out) * grad_out;
+        r_dy[ii] = grad_i;
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++) {
+        int idx = tid + ii * blockDim.x * gridDim.x;
+        if(idx < features * batch_size) {
+          dX[idx] = r_dy[ii];
+        }
+      }
+    }
+  }
+}
+
+// Compute grid size for pointwise backward kernel.
+// block_x/y is total elment being handled per block, not number of threads
+void get_biasAddRelu_bprop_grid_size(
+    int yfeat,
+    int batch_size,
+    int block_x,
+    int block_y,
+    int* grid_x,
+    int* grid_y) {
+
+  *grid_x = (yfeat + block_x - 1) / block_x;
+  // Get number of SMs for efficient reduction.
+  int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+  // can switch to occupancy calculation. use 4 below now for sm_70
+  int max_blocks_y = (num_SMs * 4+(*grid_x)-1) / (*grid_x);
+  // block_y should be from minimal work per thread
+  int nRedSplits = (batch_size + block_y - 1) / block_y;
+  // increase number of elem per thread redcution to not launch more than enough
+  // kernel adjust work, so here we just launch max block
+  *grid_y = std::min(nRedSplits, max_blocks_y);
+  return;
+}
+
+// Addition done deterministically via a 2-pass approach. Each CTA writes out partial
+// sum, and the last CTA in grid Y dimension accumulates partials serially and writes to result.
+template <typename T, int UNROLL_FACTOR>
+__global__ void biasAdd_bprop(
+    T* dY,
+    int features,
+    int batch_size,
+    volatile float* intermediate,
+    int* semaphores,
+    T* db) {
+  // The feature that this thread is responsible for
+  int f = blockIdx.x * blockDim.x + threadIdx.x;
+
+  // Compute the span this thread is responsible for
+  // For this block
+  int b_chunkSize = (batch_size + gridDim.y - 1) / gridDim.y;
+  int b_nStart = blockIdx.y * b_chunkSize;
+  int b_nSpan = min(batch_size, b_nStart + b_chunkSize) - b_nStart;
+  // For this thread
+  int chunkSize = (b_chunkSize + blockDim.y - 1) / blockDim.y;
+  int nStart = threadIdx.y * chunkSize + b_nStart;
+  int nSpan = min(b_nStart + b_nSpan, nStart + chunkSize) - nStart;
+
+  volatile float* out = intermediate + blockIdx.y * features;
+
+  // Flag to trigger last reduction.
+  __shared__ bool isLastBlock;
+  // we know block size for now
+  __shared__ float smem[BIAS_RELU_BW_NTHREADS_X*BIAS_RELU_BW_NTHREADS_Y];
+
+  // Accumulate db in FP32 always
+  float db_local = 0;
+  if (f < features) {
+    int nidx = 0;
+    // Handle non-multiple of UNROLL_FACTOR residue
+    for (; nidx < nSpan % UNROLL_FACTOR; nidx++) {
+      int64_t row, col, flat_idx;
+      row = f;
+      col = nStart + nidx;
+      flat_idx = col * features + row;
+      db_local += (float)dY[flat_idx];
+    }
+
+    // Handle meat of work
+    for (; (nidx + UNROLL_FACTOR - 1) < nSpan; nidx += UNROLL_FACTOR) {
+      int64_t row, col, flat_idx;
+      row = f;
+      col = nStart + nidx;
+      flat_idx = col * features + row;
+#pragma unroll 4
+      for (int u = 0; u < UNROLL_FACTOR; u++) {
+        db_local += (float)dY[flat_idx];
+        flat_idx += features;
+      }
+    }
+
+    // naive block reduction on y-dim
+    int linear_idx = threadIdx.y * blockDim.x + threadIdx.x;
+    smem[linear_idx] = db_local;
+  }
+  __syncthreads();
+  if (f < features) {
+    if(threadIdx.y == 0) {
+      for(int yidx = 1; yidx < blockDim.y; yidx++){
+        db_local += smem[yidx * blockDim.x + threadIdx.x];
+      }
+
+      // block result is in db_local now for all threadIdx.y == 0
+      // Write out partial result
+      out[f] = db_local;
+    }
+  }
+  __threadfence();
+  __syncthreads();
+
+  // Increment semaphore and check if this is the last CTA in the grid_y dimension.
+  // Only thread (0,0) calls this
+  if (threadIdx.x == 0 && threadIdx.y == 0 && f < features) {
+    unsigned int sum_idx;
+    sum_idx = atomicAdd(&(semaphores[blockIdx.x]), 1);
+    isLastBlock = (sum_idx == (gridDim.y - 1));
+  }
+  __syncthreads();
+
+  db_local = 0;
+  // No block reduction for now, only thread (*,0) do grid reduction
+  if (isLastBlock && f < features) {
+    if(threadIdx.y == 0) {
+      for (int n = 0; n < gridDim.y; n++) {
+        int row, col;
+        row = f;
+        col = n;
+        db_local += (float)(intermediate[col * features + row]);
+      }
+      db[f] = (T)db_local;
+    }
+  }
+}
+
+// Addition done deterministically via a 2-pass approach. Each CTA writes out partial
+// sum, and the last CTA in grid Y dimension accumulates partials serially and writes to result.
+template <typename T, int UNROLL_FACTOR>
+__global__ void biasAddRelu_bprop(
+    T* Y,
+    T* dY,
+    int features,
+    int batch_size,
+    T* dX,
+    volatile float* intermediate,
+    int* semaphores,
+    T* db) {
+  // The feature that this thread is responsible for
+  int f = blockIdx.x * blockDim.x + threadIdx.x;
+
+  // Compute the span this thread is responsible for
+  // For this block
+  int b_chunkSize = (batch_size + gridDim.y - 1) / gridDim.y;
+  int b_nStart = blockIdx.y * b_chunkSize;
+  int b_nSpan = min(batch_size, b_nStart + b_chunkSize) - b_nStart;
+  // For this thread
+  int chunkSize = (b_chunkSize + blockDim.y - 1) / blockDim.y;
+  int nStart = threadIdx.y * chunkSize + b_nStart;
+  int nSpan = min(b_nStart + b_nSpan, nStart + chunkSize) - nStart;
+
+  volatile float* out = intermediate + blockIdx.y * features;
+
+  // Flag to trigger last reduction.
+  __shared__ bool isLastBlock;
+  // we know block size for now
+  __shared__ float smem[BIAS_RELU_BW_NTHREADS_X*BIAS_RELU_BW_NTHREADS_Y];
+
+  // Accumulate db in FP32 always
+  float db_local = 0;
+  if (f < features) {
+    int nidx = 0;
+    // Handle non-multiple of UNROLL_FACTOR residue
+    for (; nidx < nSpan % UNROLL_FACTOR; nidx++) {
+      int row, col, flat_idx;
+      row = f;
+      col = nStart + nidx;
+      flat_idx = col * features + row;
+      T y_val = Y[flat_idx];
+      T dy_val = dY[flat_idx];
+      T dx_val;
+      if ((float)y_val > 0.f)
+        dx_val = dy_val;
+      else
+        dx_val = 0;
+      dX[flat_idx] = dx_val;
+      db_local += (float)dx_val;
+    }
+
+    // Handle meat of work
+    for (; (nidx + UNROLL_FACTOR - 1) < nSpan; nidx += UNROLL_FACTOR) {
+      int row, col, flat_idx;
+      row = f;
+      col = nStart + nidx;
+      flat_idx = col * features + row;
+#pragma unroll 4
+      for (int u = 0; u < UNROLL_FACTOR; u++) {
+        T y_val = Y[flat_idx];
+        T dy_val = dY[flat_idx];
+        T dx_val;
+        if ((float)y_val > 0.f)
+          dx_val = dy_val;
+        else
+          dx_val = 0;
+        dX[flat_idx] = dx_val;
+        db_local += (float)dx_val;
+        flat_idx += features;
+      }
+    }
+
+    // naive block reduction on y-dim
+    int linear_idx = threadIdx.y * blockDim.x + threadIdx.x;
+    smem[linear_idx] = db_local;
+  }
+  __syncthreads();
+  if (f < features) {
+    if(threadIdx.y == 0) {
+      for(int yidx = 1; yidx < blockDim.y; yidx++){
+        db_local += smem[yidx * blockDim.x + threadIdx.x];
+      }
+
+      // block result is in db_local now for all threadIdx.y == 0
+      // Write out partial result
+      out[f] = db_local;
+    }
+  }
+  __threadfence();
+  __syncthreads();
+
+  // Increment semaphore and check if this is the last CTA in the grid_y dimension.
+  // Only thread (0,0) calls this
+  if (threadIdx.x == 0 && threadIdx.y == 0 && f < features) {
+    unsigned int sum_idx;
+    sum_idx = atomicAdd(&(semaphores[blockIdx.x]), 1);
+    isLastBlock = (sum_idx == (gridDim.y - 1));
+  }
+  __syncthreads();
+
+  db_local = 0;
+  // No block reduction for now, only thread (*,0) do grid reduction
+  if (isLastBlock && f < features) {
+    if(threadIdx.y == 0) {
+      for (int n = 0; n < gridDim.y; n++) {
+        int row, col;
+        row = f;
+        col = n;
+        db_local += (float)(intermediate[col * features + row]);
+      }
+      db[f] = (T)db_local;
+    }
+  }
+}
+
+// Addition done deterministically via a 2-pass approach. Each CTA writes out partial
+// sum, and the last CTA in grid Y dimension accumulates partials serially and writes to result.
+template <typename T, int UNROLL_FACTOR>
+__global__ void biasAddRelu_bprop_aligned(
+    T* Y,
+    T* dY,
+    int features,
+    int batch_size,
+    T* dX,
+    volatile float* intermediate,
+    int* semaphores,
+    T* db) {
+  // The feature that this thread is responsible for
+  int f = blockIdx.x * blockDim.x + threadIdx.x;
+
+  // Compute the span this thread is responsible for
+  // For this block
+  int b_chunkSize = (batch_size + gridDim.y - 1) / gridDim.y;
+  int b_nStart = blockIdx.y * b_chunkSize;
+  int b_nSpan = min(batch_size, b_nStart + b_chunkSize) - b_nStart;
+  // For this thread
+  int chunkSize = (b_chunkSize + blockDim.y - 1) / blockDim.y;
+  int nStart = threadIdx.y * chunkSize + b_nStart;
+  int nSpan = min(b_nStart + b_nSpan, nStart + chunkSize) - nStart;
+
+  volatile float* out = intermediate + blockIdx.y * features;
+
+  // Flag to trigger last reduction.
+  __shared__ bool isLastBlock;
+
+  // Accumulate db in FP32 always
+  float db_local[ILP];
+  T r_y[ILP];
+  T r_dy[ILP];
+#pragma unroll
+  for(int ii=0;ii<ILP;ii++){
+    db_local[ii] = 0.f;
+  }
+
+  // f always <= features in this case
+  //if (f < features) {
+  int nidx = 0;
+
+  // Handle non-multiple of UNROLL_FACTOR residue
+  for (; nidx < nSpan % UNROLL_FACTOR; nidx++) {
+    int row, col, flat_idx;
+    row = f;
+    col = nStart + nidx;
+    flat_idx = col * features / ILP + row;
+
+    load_store(r_y, Y, 0, flat_idx);
+    load_store(r_dy, dY, 0, flat_idx);
+#pragma unroll
+    for(int ii=0;ii<ILP;ii++){
+      if ((float)r_y[ii] <= 0.f)
+        r_dy[ii] = 0;
+      db_local[ii] += (float)r_dy[ii];
+    }
+    load_store(dX, r_dy, flat_idx, 0);
+  }
+
+  // Handle meat of work
+  for (; (nidx + UNROLL_FACTOR - 1) < nSpan; nidx += UNROLL_FACTOR) {
+    int row, col, flat_idx;
+    row = f;
+    col = nStart + nidx;
+    flat_idx = col * features / ILP + row; // total threads in x == features/ILP
+#pragma unroll
+    for (int u = 0; u < UNROLL_FACTOR; u++) {
+      load_store(r_y, Y, 0, flat_idx);
+      load_store(r_dy, dY, 0, flat_idx);
+#pragma unroll
+      for(int ii=0;ii<ILP;ii++){
+        if ((float)r_y[ii] <= 0.f)
+          r_dy[ii] = 0;
+        db_local[ii] += (float)r_dy[ii];
+      }
+      load_store(dX, r_dy, flat_idx, 0);
+      flat_idx += features/ILP;
+    }
+  }
+
+  // we know block size for now
+  __shared__ float smem[BIAS_RELU_BW_NTHREADS_X*BIAS_RELU_BW_NTHREADS_Y*ILP];
+  // naive block reduction on y-dim
+  int linear_idx = threadIdx.y * blockDim.x + threadIdx.x;
+  float* smem_out = smem + ILP * linear_idx;
+#pragma unroll
+  for(int ii=0;ii<ILP;ii++){
+    smem_out[ii] = db_local[ii]; // reuse local dy buffer
+  }
+  __syncthreads();
+  if(threadIdx.y == 0) {
+    for(int yidx = 1; yidx < blockDim.y; yidx++){
+      float* smem_in = smem + ILP * (yidx * blockDim.x + threadIdx.x);
+#pragma unroll
+      for(int ii=0;ii<ILP;ii++){
+        db_local[ii] += smem_in[ii]; // reuse local dy buffer
+      }
+    }
+
+    // block result is in db_local now for all threadIdx.y == 0
+    if(gridDim.y == 1) {
+#pragma unroll
+      for(int ii=0;ii<ILP;ii++){
+        r_dy[ii] = db_local[ii]; // reuse local dy buffer
+      }
+      load_store(db, r_dy, f, 0);
+      return;
+    }
+
+    // Write out partial result
+    load_store(out, db_local, f, 0);
+  }
+  __threadfence();
+  __syncthreads();
+
+  // Increment semaphore and check if this is the last CTA in the grid_y dimension.
+  // Only thread (0,0) calls this
+  if (threadIdx.x == 0 && threadIdx.y == 0) {
+    unsigned int sum_idx;
+    sum_idx = atomicAdd(&(semaphores[blockIdx.x]), 1);
+    isLastBlock = (sum_idx == (gridDim.y - 1));
+  }
+  __syncthreads();
+
+#pragma unroll
+  for(int ii=0;ii<ILP;ii++){
+    db_local[ii] = 0.f;
+  }
+  float r_db[ILP];
+
+  // No block reduction for now, only thread (*,0) do grid reduction
+  if (isLastBlock) {
+    if(threadIdx.y == 0){
+      for (int n = 0; n < gridDim.y; n++) {
+        int row, col;
+        row = f;
+        col = n;
+        load_store(r_db, intermediate, 0, col * features / ILP + row);
+#pragma unroll
+        for(int ii=0;ii<ILP;ii++){
+          db_local[ii] += r_db[ii];
+        }
+      }
+#pragma unroll
+      for(int ii=0;ii<ILP;ii++){
+        r_dy[ii] = db_local[ii]; // reuse local dy buffer
+      }
+      load_store(db, r_dy, f, 0);
+    }
+  }
+}
+
+// Lists where the num_layers-1 intermediate Y buffers start in reserved space on fprop, starting
+// offset 0. The last Y value is, of course, stored in the user provided output buffer.
+void get_y_offsets(
+    int batch_size,
+    int num_layers,
+    const int* output_features,
+    int* y_start_offsets) {
+  y_start_offsets[0] = 0;
+  for (int i = 1; i < num_layers; i++) {
+    y_start_offsets[i] = y_start_offsets[i - 1] + batch_size * output_features[i - 1];
+  }
+}
+
+// Returns the reserved space (in elements) needed for the MLP
+size_t get_mlp_reserved_space(int64_t batch_size, int num_layers, const int* output_features) {
+  size_t res_space = 0;
+  // Need to store output of every intermediate MLP - size equal to output_features[i] * batch_size
+  // for all 'i' in [0, num_layers-1)
+  for (int l = 0; l < num_layers; l++) {
+    res_space += output_features[l] * batch_size;
+  }
+  return res_space;
+}
+
+// Returns the size of all fprop activations combined
+size_t get_all_activations_size(int64_t batch_size, int num_layers, const int* output_features) {
+  size_t acts_size = 0;
+  for (int l = 0; l < num_layers; l++) {
+    acts_size += output_features[l] * batch_size;
+  }
+  return acts_size;
+}
+
+#if 0
+// Returns the work space (in elements) needed for the MLP bprop.
+size_t get_mlp_bp_workspace (int batch_size, int num_layers, const int* output_features) {
+    /*
+       Workspace is partitioned as
+       DY_GEMMs : DX_GEMMs
+    */
+    size_t work_space = 0;
+
+    // Store each intermediate dY explicitly. Need 2 dYs per MLP layer (one for o/p
+    // of biasReLU_bp and one for o/p of dgrad GEMM).
+    work_space += 2*get_all_activations_size(batch_size, num_layers, output_features);
+
+    return work_space;
+}
+#endif
+
+// Scratch space needed for reductions in number of elements
+size_t get_reduction_scratch_space(int batch_size, int num_layers, const int* output_features) {
+  size_t max_scratch_space = 0;
+  // Loop over all layers to see which one needs the max scratch space
+  for (int l = 0; l < num_layers; l++) {
+    // need to find max(aligned, not_aligned)
+    int tmp, res0, res1;
+
+    int block_x = BIAS_RELU_BW_NTHREADS_X;
+    int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;
+    get_biasAddRelu_bprop_grid_size(
+      output_features[l], batch_size, block_x, block_y, &tmp, &res0);
+
+    block_x = ILP * BIAS_RELU_BW_NTHREADS_X;
+    get_biasAddRelu_bprop_grid_size(
+      output_features[l], batch_size, block_x, block_y, &tmp, &res1);
+
+    max_scratch_space = std::max(max_scratch_space, (size_t)(output_features[l] * res0));
+    max_scratch_space = std::max(max_scratch_space, (size_t)(output_features[l] * res1));
+  }
+
+  return max_scratch_space;
+}
+
+// Buffer for semaphores
+size_t get_semaphores_size(int num_layers, const int* output_features) {
+  // Upper bound on semaphores is one per feature for the layer
+  // with the most features.
+  int max_features = 0;
+  for (int l = 0; l < num_layers; l++) {
+    max_features = std::max(max_features, output_features[l]);
+  }
+  return (size_t)max_features;
+}
+
+// Returns the work space (in elements) needed for the MLP bprop.
+template <typename T>
+size_t get_mlp_bp_workspace_in_bytes(int batch_size, int num_layers, const int* output_features) {
+  size_t work_space = 0;
+
+  // Store each intermediate dY explicitly. Need 2 dYs per MLP layer (one for o/p
+  // of biasReLU_bp and one for o/p of dgrad GEMM).
+  work_space += 2 * get_all_activations_size(batch_size, num_layers, output_features) * sizeof(T);
+  work_space +=
+      get_reduction_scratch_space(batch_size, num_layers, output_features) * sizeof(float);
+  work_space += get_semaphores_size(num_layers, output_features) * sizeof(int);
+
+  return work_space;
+}
+
+// Returns pointers to each segment of the workspace
+template <typename T>
+void partition_mlp_bp_workspace(
+    int batch_size,
+    int num_layers,
+    const int* output_features,
+    void* work_space,
+    T** dy_gemms,
+    T** dx_gemms,
+    float** db_scratch,
+    int** semaphores) {
+  /*
+     Workspace is partitioned as
+     DY_GEMMs : DX_GEMMs : DB_SCRATCH : SEMAPHORES
+  */
+  // Start address where dy_gemm tensors are stored
+  *dy_gemms = reinterpret_cast<T*>(work_space);
+  // Start address where dx_gemm tensors are stored
+  *dx_gemms = *dy_gemms + get_all_activations_size(batch_size, num_layers, output_features);
+  // Start address where db intermediate tensors are stored
+  *db_scratch = reinterpret_cast<float*>(
+      *dx_gemms + get_all_activations_size(batch_size, num_layers, output_features));
+  // Start address of semaphores
+  *semaphores = reinterpret_cast<int*>(
+      *db_scratch + get_reduction_scratch_space(batch_size, num_layers, output_features));
+
+  return;
+}
+
+// Does a simple MLP fprop (GEMM+bias+ReLU).
+// Can handle num_layers number of layers, each with its own shape. Output of layer i is assumed
+// to be input of layer i+1. output_features, WPtr and BPtr are arrays of length num_layers, and
+// must be in the same order i.e. WPtr[i] and BPtr[i] are respectively the weight and bias of layer
+// 'i'.
+template <typename T>
+int mlp_fp(
+    T* X,
+    int input_features,
+    int batch_size,
+    T** WPtr,
+    int num_layers,
+    int* output_features,
+    T** BPtr,
+    T* Y,
+    T* reserved_space,
+    int use_bias,
+    int activation,
+    void* lt_workspace) {
+  T *weight, *input, *output, *bias;
+  T *reserved_space_x, *reserved_space_y;
+  reserved_space_x = NULL;
+  reserved_space_y = reserved_space;
+
+  // Get cublas handle from Pytorch
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  // Get the stream from cublas handle to reuse for biasReLU kernel.
+  cudaStream_t stream;
+  cublasGetStream(handle, &stream);
+
+  for (int layer = 0; layer < num_layers; layer++) {
+    weight = WPtr[layer];
+    input = (layer == 0) ? X : reserved_space_x;
+    output = (layer == num_layers - 1) ? Y : reserved_space_y;
+    if (use_bias) {
+      bias = BPtr[layer];
+    }
+    int ifeat = (layer == 0) ? input_features : output_features[layer - 1];
+    int ofeat = output_features[layer];
+
+    float one = 1.f;
+    float zero = 0.f;
+
+    // try with cublaslt first for supported case with valid handle
+    int cublaslt_status = 1;
+#if defined(CUBLAS_VERSION) && CUBLAS_VERSION >= 11000
+    if(activation < 1){
+        cublaslt_status = mlp_gemm_lt(
+          //ltHandle,
+          (cublasLtHandle_t)handle,
+          CUBLAS_OP_T,
+          CUBLAS_OP_N,
+          ofeat,
+          batch_size,
+          ifeat,
+          &one,
+          weight,
+          ifeat,
+          input,
+          ifeat,
+          &zero,
+          output,
+          ofeat,
+          lt_workspace,
+          1 << 22,
+          stream,
+          use_bias == 1,
+          activation == 1,
+          bias);
+    }
+#endif
+
+    // if cublaslt failed or not executed, fallback to cublas
+    if (cublaslt_status != 0) {
+      cublasStatus_t cublas_status;
+      // Call GEMM: fprop is Y = W'X
+      cublas_status = mlp_gemm(
+        handle,
+        CUBLAS_OP_T,
+        CUBLAS_OP_N,
+        ofeat,
+        batch_size,
+        ifeat,
+        &one,
+        weight,
+        ifeat,
+        input,
+        ifeat,
+        &zero,
+        output,
+        ofeat);
+
+      if (cublas_status != CUBLAS_STATUS_SUCCESS) {
+        printf("GEMM fprop failed with %d\n", cublas_status);
+        return 1;
+      }
+
+      const uint &input_size = ofeat;
+      int num_blocks = 0;
+      int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+      // Call biasReLU
+      if(use_bias == 1) {
+        if (activation == 0) { // no activation
+          cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, biasAdd_fprop<T>, BIAS_RELU_FW_NTHREADS, 0);
+          biasAdd_fprop<<<num_SMs*num_blocks, BIAS_RELU_FW_NTHREADS, 0, stream>>>(output, bias, batch_size, input_size);
+        } else if (activation == 1) { // relu
+          cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, biasAddRelu_fprop<T>, BIAS_RELU_FW_NTHREADS, 0);
+          biasAddRelu_fprop<<<num_SMs*num_blocks, BIAS_RELU_FW_NTHREADS, 0, stream>>>(output, bias, batch_size, input_size);
+        } else if (activation == 2) { // sigmoid
+          cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, biasAdd_fprop<T>, BIAS_RELU_FW_NTHREADS, 0);
+          biasAdd_fprop<<<num_SMs*num_blocks, BIAS_RELU_FW_NTHREADS, 0, stream>>>(output, bias, batch_size, input_size);
+          cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_fprop<T>, BIAS_RELU_FW_NTHREADS, 0);
+          Sigmoid_fprop<<<num_SMs*num_blocks, BIAS_RELU_FW_NTHREADS, 0, stream>>>(output, batch_size, input_size);
+        }
+      } else {
+        // don't need to do anything in case of no activation and no bias
+        if (activation == 1) { // relu
+          cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Relu_fprop<T>, BIAS_RELU_FW_NTHREADS, 0);
+          Relu_fprop<<<num_SMs*num_blocks, BIAS_RELU_FW_NTHREADS, 0, stream>>>(output, batch_size, input_size);
+        } else if (activation == 2) { // sigmoid
+          cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_fprop<T>, BIAS_RELU_FW_NTHREADS, 0);
+          Sigmoid_fprop<<<num_SMs*num_blocks, BIAS_RELU_FW_NTHREADS, 0, stream>>>(output, batch_size, input_size);
+        }
+      }
+    }
+    // Set current output as next layer input
+    reserved_space_x = reserved_space_y;
+    // Set next layer output
+    reserved_space_y += ofeat * batch_size;
+  }
+
+  return 0;
+}
+
+// Does a simple MLP bprop (GEMM+bias+ReLU).
+// Needs reserved space to come back exactly as it was populated in fprop.
+// Does dgrad and wgrad sequentially.
+template <typename T>
+int mlp_bp(
+    T* X,
+    T* Y,
+    int input_features,
+    int batch_size,
+    T** WPtr,
+    int num_layers,
+    int* output_features,
+    T* dY,
+    T* reserved_space,
+    T* work_space,
+    T* dX,
+    T** dwPtr,
+    T** dbPtr,
+    bool requires_grad,
+    int use_bias,
+    int activation) {
+  T* weight;
+  T *dweight, *dx, *dy, *dbias;
+  T *x, *y;
+
+  // Where the dx of the biasReLU (== dy of gemm) is stored. Can be thrown away
+  // after bp call.
+  T* dy_gemm_base;
+  // Where the dx after GEMM is stored.
+  T* dx_gemm_base;
+  // Where partial reduction results are stored.
+  float* db_scratch;
+  // Semaphores for reduction.
+  int* semaphores;
+
+  partition_mlp_bp_workspace<T>(
+      batch_size,
+      num_layers,
+      output_features,
+      work_space,
+      &dy_gemm_base,
+      &dx_gemm_base,
+      &db_scratch,
+      &semaphores);
+
+  size_t semaphore_size = get_semaphores_size(num_layers, output_features) * sizeof(int);
+
+  // Get cublas handle from Pytorch
+  cublasHandle_t handle = at::cuda::getCurrentCUDABlasHandle();
+  // Get the stream from cublas handle to reuse for biasReLU kernel.
+  cudaStream_t stream;
+  cublasGetStream(handle, &stream);
+
+  int* y_offsets = (int*)malloc(num_layers * sizeof(int));
+  get_y_offsets(batch_size, num_layers, output_features, y_offsets);
+
+  for (int layer = num_layers - 1; layer >= 0; layer--) {
+    weight = WPtr[layer];
+    dweight = dwPtr[layer];
+
+    // x is read from reserved space
+    x = (layer == 0) ? X : reserved_space + y_offsets[layer - 1];
+    // dx is written in workspace for all but layer==0
+    dx = (layer == 0) ? dX : dx_gemm_base + y_offsets[layer - 1];
+
+    // y is read from reserved space
+    y = (layer == num_layers - 1) ? Y : reserved_space + y_offsets[layer];
+    // dx from layer+1
+    dy = (layer == num_layers - 1) ? dY : dx_gemm_base + y_offsets[layer];
+    // dy_gemm is written to and read immediately
+    T* dy_gemm = dy_gemm_base + y_offsets[layer];
+
+    dbias = dbPtr[layer];
+    int xfeat = (layer == 0) ? input_features : output_features[layer - 1];
+    int yfeat = output_features[layer];
+
+    float one = 1.f;
+    float zero = 0.f;
+
+    if (use_bias == 1) {
+      if (activation == 0) { // no acitvation
+        // bgrad
+        dim3 block(BIAS_RELU_BW_NTHREADS_X, BIAS_RELU_BW_NTHREADS_Y);
+        int grid_x, grid_y;
+        cudaMemsetAsync(semaphores, 0, semaphore_size, stream);
+
+        int block_x = BIAS_RELU_BW_NTHREADS_X;
+        int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;
+        get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);
+        dim3 grid(grid_x, grid_y);
+        biasAdd_bprop<T, 4><<<grid, block, 0, stream>>>(
+          dy, yfeat, batch_size, db_scratch, semaphores, dbias);
+        // bypass dgrad through reset pointer
+        dy_gemm = dy;
+      } else if (activation == 1) { // relu
+        dim3 block(BIAS_RELU_BW_NTHREADS_X, BIAS_RELU_BW_NTHREADS_Y);
+        int grid_x, grid_y;
+        cudaMemsetAsync(semaphores, 0, semaphore_size, stream);
+
+        if(yfeat % (ILP * BIAS_RELU_BW_NTHREADS_X) == 0 &&
+           is_aligned(y) &&
+           is_aligned(dy) &&
+           is_aligned(dy_gemm) &&
+           is_aligned(dbias)){
+          int block_x = ILP * BIAS_RELU_BW_NTHREADS_X;
+          int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;
+          get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);
+          dim3 grid(grid_x, grid_y);
+          biasAddRelu_bprop_aligned<T, 4><<<grid, block, 0, stream>>>(
+            y, dy, yfeat, batch_size, dy_gemm, db_scratch, semaphores, dbias);
+        } else {
+          int block_x = BIAS_RELU_BW_NTHREADS_X;
+          int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;
+          get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);
+          dim3 grid(grid_x, grid_y);
+          biasAddRelu_bprop<T, 4><<<grid, block, 0, stream>>>(
+            y, dy, yfeat, batch_size, dy_gemm, db_scratch, semaphores, dbias);
+        }
+      } else if (activation == 2) { // sigmoid
+        // activation backward
+        int num_blocks = 0;
+        int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+        cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_bprop<T>, BIAS_RELU_FW_NTHREADS, 0);
+        Sigmoid_bprop<<<num_SMs*num_blocks, BIAS_RELU_FW_NTHREADS, 0, stream>>>(dy, y, batch_size, yfeat, dy_gemm);
+
+        // bgrad, from dy_gemm
+        dim3 block(BIAS_RELU_BW_NTHREADS_X, BIAS_RELU_BW_NTHREADS_Y);
+        int grid_x, grid_y;
+        cudaMemsetAsync(semaphores, 0, semaphore_size, stream);
+
+        int block_x = BIAS_RELU_BW_NTHREADS_X;
+        int block_y = BIAS_RELU_RED_PER_THREAD * BIAS_RELU_BW_NTHREADS_Y;
+        get_biasAddRelu_bprop_grid_size(yfeat, batch_size, block_x, block_y, &grid_x, &grid_y);
+        dim3 grid(grid_x, grid_y);
+        biasAdd_bprop<T, 4><<<grid, block, 0, stream>>>(
+          dy_gemm, yfeat, batch_size, db_scratch, semaphores, dbias);
+      }
+    } else { // no bias below
+      if (activation == 0) {
+        // bypass dgrad through reset pointer
+        dy_gemm = dy;
+      } else if (activation == 1) { // relu
+        int num_blocks = 0;
+        int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+        cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Relu_bprop<T>, BIAS_RELU_FW_NTHREADS, 0);
+        Relu_bprop<<<num_SMs*num_blocks, BIAS_RELU_FW_NTHREADS, 0, stream>>>(dy, y, batch_size, yfeat, dy_gemm);
+      } else if (activation == 2) { // sigmoid
+        int num_blocks = 0;
+        int num_SMs = at::cuda::getCurrentDeviceProperties()->multiProcessorCount;
+        cudaOccupancyMaxActiveBlocksPerMultiprocessor(&num_blocks, Sigmoid_bprop<T>, BIAS_RELU_FW_NTHREADS, 0);
+        Sigmoid_bprop<<<num_SMs*num_blocks, BIAS_RELU_FW_NTHREADS, 0, stream>>>(dy, y, batch_size, yfeat, dy_gemm);
+      }
+    }
+
+    cublasStatus_t cublas_status;
+    // Call GEMM dgrad
+    if (layer > 0 || requires_grad == 1) {
+      cublas_status = mlp_gemm(
+        handle,
+        CUBLAS_OP_N,
+        CUBLAS_OP_N,
+        xfeat,
+        batch_size,
+        yfeat,
+        &one,
+        weight,
+        xfeat,
+        dy_gemm,
+        yfeat,
+        &zero,
+        dx,
+        xfeat);
+
+      if (cublas_status != CUBLAS_STATUS_SUCCESS) {
+        printf("GEMM dgrad failed with %d\n", cublas_status);
+        return 1;
+      }
+    }
+
+    // Call GEMM wgrad
+    cublas_status = mlp_gemm(
+        handle,
+        CUBLAS_OP_N,
+        CUBLAS_OP_T,
+        xfeat,
+        yfeat,
+        batch_size,
+        &one,
+        x,
+        xfeat,
+        dy_gemm,
+        yfeat,
+        &zero,
+        dweight,
+        xfeat);
+
+    if (cublas_status != CUBLAS_STATUS_SUCCESS) {
+      printf("GEMM wgrad failed with %d\n", cublas_status);
+      return 1;
+    }
+  }
+
+  return 0;
+}
+
+// Instantiate for floating point types
+template int mlp_fp<float>(
+    float* X,
+    int input_features,
+    int batch_size,
+    float** WPtr,
+    int num_layers,
+    int* output_features,
+    float** BPtr,
+    float* Y,
+    float* reserved_space,
+    int use_bias,
+    int activation,
+    void* lt_workspace);
+
+template int mlp_bp<float>(
+    float* X,
+    float* Y,
+    int input_features,
+    int batch_size,
+    float** WPtr,
+    int num_layers,
+    int* output_features,
+    float* dY,
+    float* reserved_space,
+    float* work_space,
+    float* dX,
+    float** dwPtr,
+    float** dbPtr,
+    bool requires_grad,
+    int use_bias,
+    int activation);
+
+template int mlp_fp<at::Half>(
+    at::Half* X,
+    int input_features,
+    int batch_size,
+    at::Half** WPtr,
+    int num_layers,
+    int* output_features,
+    at::Half** BPtr,
+    at::Half* Y,
+    at::Half* reserved_space,
+    int use_bias,
+    int activation,
+    void* lt_workspace);
+
+template int mlp_bp<at::Half>(
+    at::Half* X,
+    at::Half* Y,
+    int input_features,
+    int batch_size,
+    at::Half** WPtr,
+    int num_layers,
+    int* output_features,
+    at::Half* dY,
+    at::Half* reserved_space,
+    at::Half* work_space,
+    at::Half* dX,
+    at::Half** dwPtr,
+    at::Half** dbPtr,
+    bool requires_grad,
+    int use_bias,
+    int activation);
+
+template int mlp_fp<double>(
+    double* X,
+    int input_features,
+    int batch_size,
+    double** WPtr,
+    int num_layers,
+    int* output_features,
+    double** BPtr,
+    double* Y,
+    double* reserved_space,
+    int use_bias,
+    int activation,
+    void* lt_workspace);
+
+template int mlp_bp<double>(
+    double* X,
+    double* Y,
+    int input_features,
+    int batch_size,
+    double** WPtr,
+    int num_layers,
+    int* output_features,
+    double* dY,
+    double* reserved_space,
+    double* work_space,
+    double* dX,
+    double** dwPtr,
+    double** dbPtr,
+    bool requires_grad,
+    int use_bias,
+    int activation);
+
+template size_t get_mlp_bp_workspace_in_bytes<float>(
+    int batch_size,
+    int num_layers,
+    const int* output_features);
+template size_t get_mlp_bp_workspace_in_bytes<at::Half>(
+    int batch_size,
+    int num_layers,
+    const int* output_features);
+template size_t get_mlp_bp_workspace_in_bytes<double>(
+    int batch_size,
+    int num_layers,
+    const int* output_features);
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_adagrad.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_adagrad.cu
new file mode 100644
index 0000000000..699681bce3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_adagrad.cu
@@ -0,0 +1,100 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "multi_tensor_apply.cuh"
+#include "type_shim.h"
+
+#define BLOCK_SIZE 1024
+#define ILP 4
+
+typedef enum {
+  ADAGRAD_MODE_0 = 0, // L2 regularization mode.
+  ADAGRAD_MODE_1 = 1, // AdamW-style weight decay.
+
+} adagradMode_t;
+
+using MATH_T = float;
+
+template <typename T> struct AdagradFunctor {
+  __device__ __forceinline__ void
+  operator()(int chunk_size, volatile int *noop_gmem, TensorListMetadata<3> &tl,
+             const float epsilon, const float lr, adagradMode_t mode,
+             const float weight_decay) {
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    T *g = (T *)tl.addresses[0][tensor_loc];
+    g += chunk_idx * chunk_size;
+
+    T *p = (T *)tl.addresses[1][tensor_loc];
+    p += chunk_idx * chunk_size;
+
+    T *h = (T *)tl.addresses[2][tensor_loc];
+    h += chunk_idx * chunk_size;
+
+    n -= chunk_idx * chunk_size;
+
+    // see note in multi_tensor_scale_kernel.cu
+    for (int i_start = 0; i_start < n && i_start < chunk_size;
+         i_start += blockDim.x * ILP) {
+      MATH_T r_g[ILP];
+      MATH_T r_p[ILP];
+      MATH_T r_h[ILP];
+#pragma unroll
+      for (int ii = 0; ii < ILP; ii++) {
+        int i = i_start + threadIdx.x + ii * blockDim.x;
+        if (i < n && i < chunk_size) {
+          r_g[ii] = g[i];
+          r_p[ii] = p[i];
+          r_h[ii] = h[i];
+        } else {
+          r_g[ii] = MATH_T(0);
+          r_p[ii] = MATH_T(0);
+          r_h[ii] = MATH_T(0);
+        }
+      }
+#pragma unroll
+      for (int ii = 0; ii < ILP; ii++) {
+        if (mode == ADAGRAD_MODE_0) { // L2
+          r_g[ii] = r_g[ii] + weight_decay * r_p[ii];
+          r_h[ii] = r_h[ii] + r_g[ii] * r_g[ii];
+          r_p[ii] = r_p[ii] - lr * (r_g[ii] / (sqrtf(r_h[ii]) + epsilon));
+        } else { // AdamW-style
+          r_h[ii] = r_h[ii] + r_g[ii] * r_g[ii];
+          r_p[ii] = r_p[ii] - lr * (r_g[ii] / (sqrtf(r_h[ii]) + epsilon) + weight_decay * r_p[ii]);
+        }
+      }
+#pragma unroll
+      for (int ii = 0; ii < ILP; ii++) {
+        int i = i_start + threadIdx.x + ii * blockDim.x;
+        if (i < n && i < chunk_size) {
+          p[i] = r_p[ii];
+          h[i] = r_h[ii];
+        }
+      }
+    }
+  }
+};
+
+void multi_tensor_adagrad_cuda(
+    int chunk_size, at::Tensor noop_flag,
+    std::vector<std::vector<at::Tensor>> tensor_lists, const float lr,
+    const float epsilon, const int mode, const float weight_decay) {
+  using namespace at;
+
+  // Assume single type across p,g,h now
+  DISPATCH_DOUBLE_FLOAT_AND_HALF(
+      tensor_lists[0][0].scalar_type(), 0, "adagrad",
+      multi_tensor_apply<3>(BLOCK_SIZE, chunk_size, noop_flag, tensor_lists,
+                            AdagradFunctor<scalar_t_0>(), epsilon, lr,
+                            (adagradMode_t)mode, weight_decay);)
+
+  AT_CUDA_CHECK(cudaGetLastError());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_adam.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_adam.cu
new file mode 100644
index 0000000000..dacbfc15fa
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_adam.cu
@@ -0,0 +1,171 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+typedef enum{
+  ADAM_MODE_0   =0, // L2 regularization mode
+  ADAM_MODE_1   =1  // Decoupled weight decay mode(AdamW)
+} adamMode_t;
+
+using MATH_T = float;
+
+template<typename T>
+struct AdamFunctor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<4>& tl,
+    const float beta1,
+    const float beta2,
+    const float beta1_correction,
+    const float beta2_correction,
+    const float epsilon,
+    const float lr,
+    adamMode_t mode,
+    const float decay)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+
+    // potentially use to pass in list of scalar
+    // int tensor_num = tl.start_tensor_this_launch + tensor_loc;
+
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    T* g = (T*)tl.addresses[0][tensor_loc];
+    g += chunk_idx*chunk_size;
+
+    T* p = (T*)tl.addresses[1][tensor_loc];
+    p += chunk_idx*chunk_size;
+
+    T* m = (T*)tl.addresses[2][tensor_loc];
+    m += chunk_idx*chunk_size;
+
+    T* v = (T*)tl.addresses[3][tensor_loc];
+    v += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    // see note in multi_tensor_scale_kernel.cu
+    for(int i_start = 0;
+            i_start < n && i_start < chunk_size;
+            i_start += blockDim.x*ILP)
+    {
+      MATH_T r_g[ILP];
+      MATH_T r_p[ILP];
+      MATH_T r_m[ILP];
+      MATH_T r_v[ILP];
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          r_g[ii] = g[i];
+          r_p[ii] = p[i];
+          r_m[ii] = m[i];
+          r_v[ii] = v[i];
+        } else {
+          r_g[ii] = MATH_T(0);
+          r_p[ii] = MATH_T(0);
+          r_m[ii] = MATH_T(0);
+          r_v[ii] = MATH_T(0);
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        if(mode == ADAM_MODE_0) { // L2
+          r_g[ii] = r_g[ii] + (decay * r_p[ii]);
+          r_m[ii] = beta1 * r_m[ii] + (1-beta1) * r_g[ii];
+          r_v[ii] = beta2 * r_v[ii] + (1-beta2) * r_g[ii] * r_g[ii];
+          MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+          MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+          MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+          MATH_T update = next_m_unbiased / denom;
+          r_p[ii] = r_p[ii] - (lr * update);
+        }
+        else { // weight decay
+          r_m[ii] = beta1 * r_m[ii] + (1-beta1) * r_g[ii];
+          r_v[ii] = beta2 * r_v[ii] + (1-beta2) * r_g[ii] * r_g[ii];
+          MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+          MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+          MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+          MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]);
+          r_p[ii] = r_p[ii] - (lr * update);
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          p[i] = r_p[ii];
+          m[i] = r_m[ii];
+          v[i] = r_v[ii];
+        }
+      }
+    }
+  }
+};
+
+void multi_tensor_adam_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  const float lr,
+  const float beta1,
+  const float beta2,
+  const float epsilon,
+  const int step,
+  const int mode,
+  const int bias_correction,
+  const float weight_decay)
+{
+  using namespace at;
+
+  // Handle bias correction mode
+  float bias_correction1 = 1.0f, bias_correction2 = 1.0f;
+  if (bias_correction == 1) {
+    bias_correction1 = 1 - std::pow(beta1, step);
+    bias_correction2 = 1 - std::pow(beta2, step);
+  }
+
+  // Assume single type across p,g,m1,m2 now
+  DISPATCH_DOUBLE_FLOAT_AND_HALF(
+    tensor_lists[0][0].scalar_type(), 0, "adam",
+    multi_tensor_apply<4>(
+      BLOCK_SIZE,
+      chunk_size,
+      noop_flag,
+      tensor_lists,
+      AdamFunctor<scalar_t_0>(),
+      beta1,
+      beta2,
+      bias_correction1,
+      bias_correction2,
+      epsilon,
+      lr,
+      (adamMode_t) mode,
+      weight_decay); )
+
+  AT_CUDA_CHECK(cudaGetLastError());
+
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_apply.cuh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_apply.cuh
new file mode 100644
index 0000000000..d06c222fd3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_apply.cuh
@@ -0,0 +1,133 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+#include <c10/cuda/CUDAGuard.h>
+#include "compat.h"
+
+#include <assert.h>
+
+// #include <iostream>
+
+// This header is the one-stop shop for all your multi-tensor apply needs.
+
+
+// TODO:  Kernel arg size limit may be <4KB for some other cards (ie Jetson)
+constexpr int depth_to_max_tensors[5] = {110, 64, 48, 36, 30};
+constexpr int depth_to_max_blocks[5] = {320, 320, 320, 320, 320};
+
+template<int n> struct TensorListMetadata
+{
+  void* addresses[n][depth_to_max_tensors[n-1]];
+  int sizes[depth_to_max_tensors[n-1]];
+  unsigned char block_to_tensor[depth_to_max_blocks[n-1]];
+  int block_to_chunk[depth_to_max_blocks[n-1]]; // I fear this needs to be a full int.
+  int start_tensor_this_launch;
+};
+
+
+template<typename T, typename U, typename... ArgTypes>
+__global__ void multi_tensor_apply_kernel(
+    int chunk_size,
+    volatile int* noop_flag,
+    T tl,
+    U callable,
+    ArgTypes... args)
+{
+  // Hand the chunk information to the user-supplied functor to process however it likes.
+  callable(chunk_size, noop_flag, tl, args...);
+}
+
+template<int depth, typename T, typename... ArgTypes>
+void multi_tensor_apply(
+  int block_size,
+  int chunk_size,
+  const at::Tensor& noop_flag,
+  const std::vector<std::vector<at::Tensor>>& tensor_lists,
+  T callable,
+  ArgTypes... args)
+{
+  TORCH_CHECK(tensor_lists.size() == depth, "tensor_lists.size() != depth");
+  int len0 = tensor_lists[0].size();
+  TORCH_CHECK(len0 > 0, "tensor_lists[0].size() is not > 0");
+  auto ref_device = tensor_lists[0][0].device();
+  TORCH_CHECK(ref_device.type() == at::kCUDA, "expected input to be on cuda");
+  for (int l = 0; l < tensor_lists.size(); l++) // No range-based for because I need indices
+  {
+    TORCH_CHECK(tensor_lists[l].size() == len0, "Size mismatch among tensor lists");
+    for(int t = 0; t < tensor_lists[l].size(); t++)
+    {
+      // TODO:  Print which tensor fails.
+      bool contiguous_memory = tensor_lists[l][t].is_contiguous();
+#ifdef VERSION_GE_1_5
+      contiguous_memory = (contiguous_memory || tensor_lists[l][t].is_contiguous(at::MemoryFormat::ChannelsLast));
+#endif
+      TORCH_CHECK(contiguous_memory, "A tensor was not contiguous.");
+      TORCH_CHECK(tensor_lists[l][t].device() == ref_device, "A tensor was not on the same device as the first tensor");
+      TORCH_CHECK(tensor_lists[l][t].numel() == tensor_lists[0][t].numel(), "Size mismatch");
+    }
+  }
+
+  int ntensors = tensor_lists[0].size();
+
+  TensorListMetadata<depth> tl;
+
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(tensor_lists[0][0]));
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  tl.start_tensor_this_launch = 0;
+  int loc_block_info = 0;
+  int loc_tensor_info = 0;
+  for(int t = 0; t < ntensors; t++)
+  {
+    tl.sizes[loc_tensor_info] = tensor_lists[0][t].numel();
+    for(int d = 0; d < depth; d++)
+      tl.addresses[d][loc_tensor_info] = tensor_lists[d][t].data_ptr();
+    loc_tensor_info++;
+
+    int chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1)/chunk_size;
+
+    for(int chunk = 0; chunk < chunks_this_tensor; chunk++)
+    {
+      // std::cout << chunks_this_tensor << std::endl;
+      tl.block_to_tensor[loc_block_info] = loc_tensor_info - 1;
+      tl.block_to_chunk[loc_block_info] = chunk;
+      loc_block_info++;
+
+      bool tensors_full = (loc_tensor_info == depth_to_max_tensors[depth-1] &&
+                           chunk == chunks_this_tensor - 1);
+      bool blocks_full = (loc_block_info == depth_to_max_blocks[depth-1]);
+      bool last_chunk = (t == ntensors - 1 && chunk == chunks_this_tensor - 1);
+      if(tensors_full || blocks_full || last_chunk)
+      {
+        // using accscalar_t = acc_type<scalar_t, true>;
+        multi_tensor_apply_kernel<<<loc_block_info, block_size, 0, stream>>>(
+          chunk_size,
+          noop_flag.DATA_PTR<int>(),
+          tl,
+          callable,
+          args...);
+
+        AT_CUDA_CHECK(cudaGetLastError());
+
+        // Reset.  The control flow possibilities here make my brain hurt.
+        loc_block_info = 0;
+        if(chunk == chunks_this_tensor - 1)
+        {
+          // std::cout << "Hit case 1 " << cond1 << " " << cond2 << " " << cond3 << std::endl;
+          loc_tensor_info = 0;
+          tl.start_tensor_this_launch = t + 1;
+        }
+        else
+        {
+          // std::cout << "Hit case 2 " << cond1 << " " << cond2 << " " << cond3 << std::endl;
+          tl.sizes[0] = tl.sizes[loc_tensor_info-1];
+          for(int d = 0; d < depth; d++)
+            tl.addresses[d][0] = tl.addresses[d][loc_tensor_info-1];
+          loc_tensor_info = 1;
+          tl.start_tensor_this_launch = t;
+        }
+      }
+    }
+  }
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_axpby_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_axpby_kernel.cu
new file mode 100644
index 0000000000..021df27d70
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_axpby_kernel.cu
@@ -0,0 +1,157 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+template<typename T>
+__device__ __forceinline__ bool is_aligned(T* p){
+  return ((uint64_t)p) % (ILP*sizeof(T)) == 0;
+}
+
+template<typename T>
+__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+
+template<typename x_t, typename y_t, typename out_t>
+struct AxpbyFunctor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<3>& tl,
+    float a,
+    float b,
+    int arg_to_check)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    x_t* x = (x_t*)tl.addresses[0][tensor_loc];
+    x += chunk_idx*chunk_size;
+
+    y_t* y = (y_t*)tl.addresses[1][tensor_loc];
+    y += chunk_idx*chunk_size;
+
+    out_t* out = (out_t*)tl.addresses[2][tensor_loc];
+    out += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    bool finite = true;
+    x_t r_x[ILP];
+    y_t r_y[ILP];
+    out_t r_out[ILP];
+
+    // to make things simple, we put aligned case in a different code path
+    if(n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x) && is_aligned(y) && is_aligned(out))
+    {
+      for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+      {
+        // load
+        load_store(r_x, x, 0 , i_start);
+        load_store(r_y, y, 0 , i_start);
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_out[ii] = a*static_cast<float>(r_x[ii]) + b*static_cast<float>(r_y[ii]);
+          if(arg_to_check == -1)
+            finite = finite && (isfinite(r_x[ii]) && isfinite(r_y[ii]));
+          if(arg_to_check == 0)
+            finite = finite && isfinite(r_x[ii]);
+          if(arg_to_check == 1)
+            finite = finite && isfinite(r_y[ii]);
+        }
+        // store
+        load_store(out, r_out, i_start , 0);
+      }
+    }
+    else
+    {
+      // Non-divergent exit condition for __syncthreads, not necessary here
+      for(int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x*ILP)
+      {
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_x[ii] = 0;
+          r_y[ii] = 0;
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            r_x[ii] = x[i];
+            r_y[ii] = y[i];
+          }
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_out[ii] = a*static_cast<float>(r_x[ii]) + b*static_cast<float>(r_y[ii]);
+          if(arg_to_check == -1)
+            finite = finite && (isfinite(r_x[ii]) && isfinite(r_y[ii]));
+          if(arg_to_check == 0)
+            finite = finite && isfinite(r_x[ii]);
+          if(arg_to_check == 1)
+            finite = finite && isfinite(r_y[ii]);
+        }
+        // see note in multi_tensor_scale_kernel.cu
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+            out[i] = r_out[ii];
+        }
+      }
+    }
+    if(!finite)
+      *noop_gmem = 1; // Blindly fire off a write.  These will race but that's ok.
+  }
+};
+
+void multi_tensor_axpby_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  float a,
+  float b,
+  int arg_to_check)
+{
+  using namespace at;
+  // The output (downscaled) type is always float.
+  // If build times suffer, think about where to put this dispatch,
+  // and what logic should be moved out of multi_tensor_apply.
+
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "multi_tensor_axpby_cuda",
+    DISPATCH_FLOAT_AND_HALF(tensor_lists[1][0].scalar_type(), 1, "multi_tensor_axpby_cuda",
+      DISPATCH_FLOAT_AND_HALF(tensor_lists[2][0].scalar_type(), 2, "multi_tensor_axpby_cuda",
+           multi_tensor_apply<3>(
+             BLOCK_SIZE,
+             chunk_size,
+             noop_flag,
+             tensor_lists,
+             AxpbyFunctor<scalar_t_0, scalar_t_1, scalar_t_2>(),
+             a,
+             b,
+             arg_to_check); )))
+
+  AT_CUDA_CHECK(cudaGetLastError());
+
+  // AT_CUDA_CHECK(cudaDeviceSynchronize());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_l2norm_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_l2norm_kernel.cu
new file mode 100644
index 0000000000..2bf615e949
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_l2norm_kernel.cu
@@ -0,0 +1,448 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+#include <c10/cuda/CUDAGuard.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+template<typename T>
+__device__ __forceinline__ bool is_aligned(T* p){
+  return ((uint64_t)p) % (ILP*sizeof(T)) == 0;
+}
+
+template<typename T>
+__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+
+template<typename x_t>
+struct L2NormFunctor
+{
+  __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<1>& tl,
+    float* output,
+    float* output_per_tensor,
+    bool per_tensor,
+    int max_chunks_per_tensor)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    x_t* x = (x_t*)tl.addresses[0][tensor_loc];
+    x += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    __shared__ float s_vals[512];
+
+    float vals[ILP]; // = {0}; // this probably works too but I want to be sure...
+    x_t r_x[ILP];
+    for(int i = 0; i < ILP; i++)
+    {
+      vals[i] = 0.f;
+      r_x[i] = 0;
+    }
+
+    // to make things simple, we put aligned case in a different code path
+    if(n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x))
+    {
+      for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+      {
+        // load
+        load_store(r_x, x, 0 , i_start);
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          float next = static_cast<float>(r_x[ii]);
+          vals[ii] += next*next;
+        }
+      }
+    }
+    else
+    {
+      for(int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x*ILP)
+      {
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            float next = static_cast<float>(x[i]);
+            vals[ii] += next*next;
+          }
+        }
+      }
+    }
+
+    float val = 0.f;
+    for(int i = 0; i < ILP; i++)
+        val += vals[i];
+
+    float final = reduce_block_into_lanes(s_vals, val);
+
+    if(threadIdx.x == 0)
+    {
+      if(!isfinite(final))
+        *noop_gmem = 1; // Blindly fire off a write.  These will race but that's ok.
+      output[blockIdx.x] += final;
+      if(per_tensor)
+        output_per_tensor[(tl.start_tensor_this_launch + tensor_loc)*max_chunks_per_tensor + chunk_idx] = final;
+    }
+  }
+};
+
+// Probably better to template, but since we are not likely to support other norm
+template<typename x_t>
+struct MaxNormFunctor
+{
+  __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<1>& tl,
+    float* output,
+    float* output_per_tensor,
+    bool per_tensor,
+    int max_chunks_per_tensor)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    x_t* x = (x_t*)tl.addresses[0][tensor_loc];
+    x += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    __shared__ float s_vals[512];
+
+    float vals[ILP]; // = {0}; // this probably works too but I want to be sure...
+    x_t r_x[ILP];
+    for(int i = 0; i < ILP; i++)
+    {
+      vals[i] = 0.f;
+      r_x[i] = 0;
+    }
+
+    // to make things simple, we put aligned case in a different code path
+    if(n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x))
+    {
+      for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+      {
+        // load
+        load_store(r_x, x, 0 , i_start);
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          float next = static_cast<float>(r_x[ii]);
+          vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));
+        }
+      }
+    }
+    else
+    {
+      for(int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x*ILP)
+      {
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            float next = static_cast<float>(x[i]);
+            vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));
+          }
+        }
+      }
+    }
+
+    float val = 0.f;
+    for(int i = 0; i < ILP; i++)
+        val = fmaxf(fabsf(val), fabsf(vals[i]));
+
+    float final = reduce_block_into_lanes_max_op(s_vals, val);
+
+    if(threadIdx.x == 0)
+    {
+      if(!isfinite(final))
+        *noop_gmem = 1; // Blindly fire off a write.  These will race but that's ok.
+      output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));
+      if(per_tensor)
+        output_per_tensor[(tl.start_tensor_this_launch + tensor_loc)*max_chunks_per_tensor + chunk_idx] = final;
+    }
+  }
+};
+
+
+__global__ void cleanup(
+  float* output,
+  float* output_per_tensor,
+  float* ret,
+  float* ret_per_tensor,
+  bool per_tensor,
+  int max_chunks_per_tensor)
+{
+  __shared__ float vals[512];
+
+  if(blockIdx.x == 0)
+  {
+    float val = 0;
+    if(threadIdx.x < 320)
+      val = output[threadIdx.x];
+
+    float final = reduce_block_into_lanes(vals, val);
+
+    if(threadIdx.x == 0)
+      *ret = sqrt(final);
+  }
+
+  if(per_tensor)
+  {
+    float* output_this_tensor = output_per_tensor + blockIdx.x*max_chunks_per_tensor;
+
+    float val = 0;
+    for(int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)
+      val += output_this_tensor[i];
+
+    float final = reduce_block_into_lanes(vals, val);
+
+    if(threadIdx.x == 0)
+      ret_per_tensor[blockIdx.x] = sqrt(final);
+  }
+}
+
+__global__ void cleanup_v2(
+  float* output,
+  float* output_per_tensor,
+  float* ret,
+  float* ret_per_tensor,
+  bool per_tensor,
+  int max_chunks_per_tensor,
+  int norm_type,
+  float alpha,
+  float beta)
+{
+  __shared__ float vals[512];
+
+  if(blockIdx.x == 0)
+  {
+    float val = 0;
+    if(threadIdx.x < 320)
+      val = output[threadIdx.x];
+
+    if (norm_type == 0) {
+      float final = reduce_block_into_lanes_max_op(vals, val);
+      if(threadIdx.x == 0)
+        *ret = alpha * (*ret) + beta * final;
+    }
+    else {
+      float final = reduce_block_into_lanes(vals, val);
+      if(threadIdx.x == 0)
+        *ret = sqrt(alpha * (*ret) * (*ret) + beta * final);
+    }
+  }
+
+  if(per_tensor)
+  {
+    float* output_this_tensor = output_per_tensor + blockIdx.x*max_chunks_per_tensor;
+
+    if (norm_type == 0) {
+      float val = 0;
+      for(int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)
+        val = fmaxf(fabsf(val), fabsf(output_this_tensor[i]));
+
+      float final = reduce_block_into_lanes_max_op(vals, val);
+
+      if(threadIdx.x == 0)
+        ret_per_tensor[blockIdx.x] = alpha * ret_per_tensor[blockIdx.x] + beta * final;
+    }
+    else {
+      float val = 0;
+      for(int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)
+        val += output_this_tensor[i];
+
+      float final = reduce_block_into_lanes(vals, val);
+
+      if(threadIdx.x == 0)
+        ret_per_tensor[blockIdx.x] = sqrt(alpha * ret_per_tensor[blockIdx.x] * ret_per_tensor[blockIdx.x] + beta * final);
+    }
+  }
+}
+
+std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::optional<bool> per_tensor_python)
+{
+  bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;
+
+  auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);
+  auto output = at::zeros({320}, float_options);
+
+  at::Tensor output_per_tensor;
+  at::Tensor ret_per_tensor;
+
+  int ntensors = tensor_lists[0].size();
+  int max_chunks_per_tensor = -1;
+
+  if(per_tensor)
+  {
+    for(int t = 0; t < ntensors; t++)
+    {
+      int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1)/chunk_size;
+      if(max_chunks_this_tensor > max_chunks_per_tensor)
+        max_chunks_per_tensor = max_chunks_this_tensor;
+    }
+    output_per_tensor = at::zeros({ntensors*max_chunks_per_tensor}, float_options);
+    ret_per_tensor = at::empty({ntensors}, float_options);
+  }
+  else
+  {
+    ret_per_tensor = at::empty({0}, float_options);
+  }
+
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda",
+    multi_tensor_apply<1>(
+      BLOCK_SIZE,
+      chunk_size,
+      noop_flag,
+      tensor_lists,
+      L2NormFunctor<scalar_t_0>(),
+      output.DATA_PTR<float>(),
+      per_tensor ? output_per_tensor.DATA_PTR<float>() : nullptr,
+      per_tensor,
+      max_chunks_per_tensor);)
+
+  AT_CUDA_CHECK(cudaGetLastError());
+  // AT_CUDA_CHECK(cudaDeviceSynchronize());
+
+  // This involves one more small kernel launches, but will be negligible end to end.
+  // I could get rid of these by hacking the functor + multi tensor harness with persistence
+  // logic, but keeping it simple for now
+  auto ret = at::empty({1}, output.options());
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  cleanup<<<per_tensor ? ntensors : 1, 512, 0, stream>>>(
+    output.DATA_PTR<float>(),
+    per_tensor ? output_per_tensor.DATA_PTR<float>() : nullptr,
+    ret.DATA_PTR<float>(),
+    per_tensor ? ret_per_tensor.DATA_PTR<float>() : nullptr,
+    per_tensor,
+    max_chunks_per_tensor);
+
+  return std::tuple<at::Tensor, at::Tensor>(ret, ret_per_tensor);
+}
+
+
+// Compute and update grad norm
+// Here use a per tensor norm, and blend new norm(n) and old norm(gn) by
+// L-2: gn = sqrt(a * gn^2 + b * n^2)
+// L-inf: gn = a * gn + b * n
+void multi_tensor_norm_out_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor out,
+  const float alpha,
+  const float beta,
+  const int norm_type)
+{
+  auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);
+  TORCH_CHECK(tensor_lists[0][0].device() == noop_flag.device(), "noop flag should be on the same device as tensors");
+  // we don't need global thus uses empty here
+  auto output = at::empty({320}, float_options);
+
+  at::Tensor output_per_tensor;
+  at::Tensor ret_per_tensor;
+
+  int ntensors = tensor_lists[0].size();
+  int max_chunks_per_tensor = -1;
+
+  for(int t = 0; t < ntensors; t++)
+  {
+    int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1)/chunk_size;
+    if(max_chunks_this_tensor > max_chunks_per_tensor)
+      max_chunks_per_tensor = max_chunks_this_tensor;
+  }
+
+  // Although it is single write then read, still need to be zero
+  // Since tailing element also participate cleanup
+  output_per_tensor = at::zeros({ntensors*max_chunks_per_tensor}, float_options);
+
+  if (norm_type == 0) {
+    DISPATCH_FLOAT_AND_HALF(
+      tensor_lists[0][0].scalar_type(), 0, "multi_tensor_maxnorm_cuda",
+      multi_tensor_apply<1>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        MaxNormFunctor<scalar_t_0>(),
+        output.DATA_PTR<float>(),
+        output_per_tensor.DATA_PTR<float>(),
+        true,
+        max_chunks_per_tensor);)
+  }
+  else {
+    DISPATCH_FLOAT_AND_HALF(
+      tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_cuda",
+      multi_tensor_apply<1>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        L2NormFunctor<scalar_t_0>(),
+        output.DATA_PTR<float>(),
+        output_per_tensor.DATA_PTR<float>(),
+        true,
+        max_chunks_per_tensor);)
+  }
+  AT_CUDA_CHECK(cudaGetLastError());
+
+  // AT_CUDA_CHECK(cudaDeviceSynchronize());
+
+  // This involves one more small kernel launches, but will be negligible end to end.
+  // I could get rid of these by hacking the functor + multi tensor harness with persistence
+  // logic, but keeping it simple for now
+  auto ret = at::empty({1}, output.options());
+
+  // Adding the following device guard since it happens sometimes that the 
+  // tensors are on one device and the cuda stream is on another device which  
+  // results in ILLEGAL MEM ACCESS error. 
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  cleanup_v2<<<ntensors, 512, 0, stream>>>(
+    output.DATA_PTR<float>(),
+    output_per_tensor.DATA_PTR<float>(),
+    ret.DATA_PTR<float>(),
+    out.DATA_PTR<float>(),
+    true,
+    max_chunks_per_tensor,
+    norm_type,
+    alpha,
+    beta);
+
+  return ;
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_l2norm_scale_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_l2norm_scale_kernel.cu
new file mode 100644
index 0000000000..f60e96090e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_l2norm_scale_kernel.cu
@@ -0,0 +1,326 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+#include <c10/cuda/CUDAGuard.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+template<typename T>
+__device__ __forceinline__ bool is_aligned(T* p){
+  return ((uint64_t)p) % (ILP*sizeof(T)) == 0;
+}
+
+template<typename T>
+__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+
+template<typename in_t, typename out_t>
+struct L2NormScaleFunctor
+{
+  __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<2>& tl,
+    float* output,
+    float* output_per_tensor,
+    float scale,
+    bool per_tensor,
+    int max_chunks_per_tensor)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    in_t* in = (in_t*)tl.addresses[0][tensor_loc];
+    in += chunk_idx*chunk_size;
+
+    out_t* out = (out_t*)tl.addresses[1][tensor_loc];
+    out += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    __shared__ float s_vals[512];
+
+    float vals[ILP]; // = {0}; // this probably works too but I want to be sure...
+    in_t r_in[ILP];
+    for(int i = 0; i < ILP; i++)
+    {
+      vals[i] = 0.f;
+      r_in[i] = 0;
+    }
+    //bool finite = true;
+    out_t r_out[ILP];
+
+    // to make things simple, we put aligned case in a different code path
+    if(n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(in) && is_aligned(out))
+    {
+      for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+      {
+        // load
+        load_store(r_in, in, 0 , i_start);
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          float next = static_cast<float>(r_in[ii]);
+          r_out[ii] = next*scale;
+          vals[ii] += next*next;
+          //finite = finite && isfinite(r_in[ii]);
+        }
+        load_store(out, r_out, i_start, 0);
+      }
+    }
+    else
+    {
+      for(int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x*ILP)
+      {
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_in[ii] = 0;
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            r_in[ii] = in[i];
+            float next = static_cast<float>(in[i]);
+            vals[ii] += next*next;
+          }
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_out[ii] = static_cast<float>(r_in[ii]) * scale;
+         // finite = finite && isfinite(r_in[ii]);
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+            out[i] = r_out[ii];
+        }
+      }
+    }
+
+    float val = 0.f;
+    for(int i = 0; i < ILP; i++)
+        val += vals[i];
+
+    float final = reduce_block_into_lanes(s_vals, val);
+
+    if(threadIdx.x == 0)
+    {
+      if(!isfinite(final))
+        *noop_gmem = 1; // Blindly fire off a write.  These will race but that's ok.
+      output[blockIdx.x] += final;
+      if(per_tensor)
+        output_per_tensor[(tl.start_tensor_this_launch + tensor_loc)*max_chunks_per_tensor + chunk_idx] = final;
+    }
+  }
+};
+// Probably better to template, but since we are not likely to support other norm
+template<typename x_t>
+struct MaxNormFunctor
+{
+  __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<1>& tl,
+    float* output,
+    float* output_per_tensor,
+    bool per_tensor,
+    int max_chunks_per_tensor)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    x_t* x = (x_t*)tl.addresses[0][tensor_loc];
+    x += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    __shared__ float s_vals[512];
+
+    float vals[ILP]; // = {0}; // this probably works too but I want to be sure...
+    x_t r_x[ILP];
+    for(int i = 0; i < ILP; i++)
+    {
+      vals[i] = 0.f;
+      r_x[i] = 0;
+    }
+
+    // to make things simple, we put aligned case in a different code path
+    if(n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(x))
+    {
+      for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+      {
+        // load
+        load_store(r_x, x, 0 , i_start);
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          float next = static_cast<float>(r_x[ii]);
+          vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));
+        }
+      }
+    }
+    else
+    {
+      for(int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x*ILP)
+      {
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            float next = static_cast<float>(x[i]);
+            vals[ii] = fmaxf(fabsf(vals[ii]), fabsf(next));
+          }
+        }
+      }
+    }
+
+    float val = 0.f;
+    for(int i = 0; i < ILP; i++)
+        val = fmaxf(fabsf(val), fabsf(vals[i]));
+
+    float final = reduce_block_into_lanes_max_op(s_vals, val);
+
+    if(threadIdx.x == 0)
+    {
+      if(!isfinite(final))
+        *noop_gmem = 1; // Blindly fire off a write.  These will race but that's ok.
+      output[blockIdx.x] = fmaxf(fabsf(output[blockIdx.x]), fabsf(final));
+      if(per_tensor)
+        output_per_tensor[(tl.start_tensor_this_launch + tensor_loc)*max_chunks_per_tensor + chunk_idx] = final;
+    }
+  }
+};
+
+__global__ void cleanup_v3(
+  float* output,
+  float* output_per_tensor,
+  float* ret,
+  float* ret_per_tensor,
+  bool per_tensor,
+  int max_chunks_per_tensor)
+{
+  __shared__ float vals[512];
+
+  if(blockIdx.x == 0)
+  {
+    float val = 0;
+    if(threadIdx.x < 320)
+      val = output[threadIdx.x];
+
+    float final = reduce_block_into_lanes(vals, val);
+
+    if(threadIdx.x == 0)
+      *ret = sqrt(final);
+  }
+
+  if(per_tensor)
+  {
+    float* output_this_tensor = output_per_tensor + blockIdx.x*max_chunks_per_tensor;
+
+    float val = 0;
+    for(int i = threadIdx.x; i < max_chunks_per_tensor; i += blockDim.x)
+      val += output_this_tensor[i];
+
+    float final = reduce_block_into_lanes(vals, val);
+
+    if(threadIdx.x == 0)
+      ret_per_tensor[blockIdx.x] = sqrt(final);
+  }
+}
+
+
+std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_scale_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  float scale,
+  at::optional<bool> per_tensor_python)
+{
+  bool per_tensor = per_tensor_python.has_value() ? per_tensor_python.value() : false;
+
+  auto float_options = tensor_lists[0][0].options().dtype(at::kFloat);
+  auto output = at::zeros({320}, float_options);
+
+  at::Tensor output_per_tensor;
+  at::Tensor ret_per_tensor;
+
+  int ntensors = tensor_lists[0].size();
+  int max_chunks_per_tensor = -1;
+
+  if(per_tensor)
+  {
+    for(int t = 0; t < ntensors; t++)
+    {
+      int max_chunks_this_tensor = (tensor_lists[0][t].numel() + chunk_size - 1)/chunk_size;
+      if(max_chunks_this_tensor > max_chunks_per_tensor)
+        max_chunks_per_tensor = max_chunks_this_tensor;
+    }
+    output_per_tensor = at::zeros({ntensors*max_chunks_per_tensor}, float_options);
+    ret_per_tensor = at::empty({ntensors}, float_options);
+  }
+  else
+  {
+    ret_per_tensor = at::empty({0}, float_options);
+  }
+
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "multi_tensor_l2norm_scale_cuda",
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[1][0].scalar_type(), 1, "multi_tensor_l2norm_scale_cuda",
+    multi_tensor_apply<2>(
+      BLOCK_SIZE,
+      chunk_size,
+      noop_flag,
+      tensor_lists,
+      L2NormScaleFunctor<scalar_t_0, scalar_t_1>(),
+      output.DATA_PTR<float>(),
+      per_tensor ? output_per_tensor.DATA_PTR<float>() : nullptr,
+      scale,
+      per_tensor,
+      max_chunks_per_tensor);))
+
+  AT_CUDA_CHECK(cudaGetLastError());
+  // AT_CUDA_CHECK(cudaDeviceSynchronize());
+
+  // This involves one more small kernel launches, but will be negligible end to end.
+  // I could get rid of these by hacking the functor + multi tensor harness with persistence
+  // logic, but keeping it simple for now
+  auto ret = at::empty({1}, output.options());
+  const at::cuda::OptionalCUDAGuard device_guard(device_of(output));
+  auto stream = at::cuda::getCurrentCUDAStream();
+  cleanup_v3<<<per_tensor ? ntensors : 1, 512, 0, stream>>>(
+    output.DATA_PTR<float>(),
+    per_tensor ? output_per_tensor.DATA_PTR<float>() : nullptr,
+    ret.DATA_PTR<float>(),
+    per_tensor ? ret_per_tensor.DATA_PTR<float>() : nullptr,
+    per_tensor,
+    max_chunks_per_tensor);
+
+  return std::tuple<at::Tensor, at::Tensor>(ret, ret_per_tensor);
+}
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb.cu
new file mode 100644
index 0000000000..3137fcd215
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb.cu
@@ -0,0 +1,413 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+template<typename T>
+__device__ __forceinline__ bool is_aligned(T* p){
+  return ((uint64_t)p) % (ILP*sizeof(T)) == 0;
+}
+
+template<typename T>
+__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+
+typedef enum{
+  MOMENT_MODE_0   =0, // L2 regularization mode
+  MOMENT_MODE_1   =1  // Decoupled weight decay mode
+} adamMode_t;
+
+std::tuple<at::Tensor, at::Tensor> multi_tensor_l2norm_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::optional<bool> per_tensor_python);
+
+using MATH_T = float;
+
+template<typename T>
+struct LAMBStage1Functor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<4>& tl,
+    const float beta1,
+    const float beta2,
+    const float beta3,
+    const float beta1_correction,
+    const float beta2_correction,
+    const float epsilon,
+    adamMode_t mode,
+    const float decay,
+    const float* global_grad_norm,
+    const float max_global_grad_norm)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    float clipped_global_grad_norm = (*global_grad_norm) > max_global_grad_norm ? (*global_grad_norm) / max_global_grad_norm : 1.0f;
+
+    T* g = (T*)tl.addresses[0][tensor_loc];
+    g += chunk_idx*chunk_size;
+
+    T* p = (T*)tl.addresses[1][tensor_loc];
+    p += chunk_idx*chunk_size;
+
+    T* m = (T*)tl.addresses[2][tensor_loc];
+    m += chunk_idx*chunk_size;
+
+    T* v = (T*)tl.addresses[3][tensor_loc];
+    v += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    MATH_T r_g[ILP];
+    MATH_T r_p[ILP];
+    MATH_T r_m[ILP];
+    MATH_T r_v[ILP];
+    // to make things simple, we put aligned case in a different code path
+    if(n % ILP == 0 &&
+       chunk_size % ILP == 0 &&
+       is_aligned(g) &&
+       is_aligned(p) &&
+       is_aligned(m) &&
+       is_aligned(v))
+    {
+      T l_g[ILP];
+      T l_p[ILP];
+      T l_m[ILP];
+      T l_v[ILP];
+      for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+      {
+        // load
+        load_store(l_g, g, 0, i_start);
+        if (decay != 0)
+          load_store(l_p, p, 0, i_start);
+        load_store(l_m, m, 0, i_start);
+        load_store(l_v, v, 0, i_start);
+        // unpack
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_g[ii] = l_g[ii];
+          if (decay == 0) {
+            r_p[ii] = MATH_T(0);
+          }
+          else {
+            r_p[ii] = l_p[ii];
+          }
+          r_m[ii] = l_m[ii];
+          r_v[ii] = l_v[ii];
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          if (mode == MOMENT_MODE_0) {
+            MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
+            // L2 on scaled grad
+            scaled_grad = scaled_grad + decay*r_p[ii];
+            r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
+            r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+            MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+            MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+            MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+            r_p[ii] = next_m_unbiased / denom;
+          }
+          else {
+            MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
+            r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
+            r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+            MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+            MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+            MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+            r_p[ii] = (next_m_unbiased/denom) + (decay*r_p[ii]);
+          }
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          l_p[ii] = r_p[ii];
+          l_m[ii] = r_m[ii];
+          l_v[ii] = r_v[ii];
+        }
+        // store
+        load_store(g, l_p, i_start, 0);
+        load_store(m, l_m, i_start, 0);
+        load_store(v, l_v, i_start, 0);
+      }
+    }
+    else
+    {
+      // see note in multi_tensor_scale_kernel.cu
+      for(int i_start = 0;
+          i_start < n && i_start < chunk_size;
+          i_start += blockDim.x*ILP)
+      {
+        MATH_T r_g[ILP];
+        MATH_T r_p[ILP];
+        MATH_T r_m[ILP];
+        MATH_T r_v[ILP];
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            r_g[ii] = g[i];
+            // special ?optimization? for lamb stage 1
+            if (decay == 0) {
+              r_p[ii] = MATH_T(0);
+            }
+            else {
+              r_p[ii] = p[i];
+            }
+            r_m[ii] = m[i];
+            r_v[ii] = v[i];
+          } else {
+            r_g[ii] = MATH_T(0);
+            r_p[ii] = MATH_T(0);
+            r_m[ii] = MATH_T(0);
+            r_v[ii] = MATH_T(0);
+          }
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          if (mode == MOMENT_MODE_0) {
+            MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
+            // L2 on scaled grad
+            scaled_grad = scaled_grad + decay*r_p[ii];
+            r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
+            r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+            MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+            MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+            MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+            r_p[ii] = next_m_unbiased / denom;
+          }
+          else {
+            MATH_T scaled_grad = r_g[ii] / clipped_global_grad_norm;
+            r_m[ii] = r_m[ii] * beta1 + beta3 * scaled_grad;
+            r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+            MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+            MATH_T next_v_unbiased = r_v[ii] / beta2_correction;
+            MATH_T denom = sqrtf(next_v_unbiased) + epsilon;
+            r_p[ii] = (next_m_unbiased/denom) + (decay*r_p[ii]);
+          }
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            g[i] = r_p[ii];
+            m[i] = r_m[ii];
+            v[i] = r_v[ii];
+          }
+        }
+      }
+    }
+  }
+};
+
+// Step 2 reads in 'update' value and per-tensor param_norm and update_norm.
+// It computes new parameter value.
+template<typename T>
+struct LAMBStage2Functor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<2>& tl,
+    const float* per_tensor_param_norm,
+    const float* per_tensor_update_norm,
+    const float learning_rate,
+    const float decay,
+    bool use_nvlamb)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int tensor_num = tl.start_tensor_this_launch + tensor_loc;
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    MATH_T ratio = learning_rate;
+    // nvlamb: apply adaptive learning rate to all parameters
+    // otherwise, only apply to those with non-zero weight decay
+    if (use_nvlamb || (decay != 0.0))
+    {
+      float param_norm = per_tensor_param_norm[tensor_num];
+      float update_norm = per_tensor_update_norm[tensor_num];
+      ratio = (update_norm != 0.0f && param_norm != 0.0f) ? learning_rate * (param_norm / update_norm) : learning_rate;
+    }
+
+    T* update = (T*)tl.addresses[0][tensor_loc];
+    update += chunk_idx*chunk_size;
+
+    T* p = (T*)tl.addresses[1][tensor_loc];
+    p += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    // to make things simple, we put aligned case in a different code path
+    if(n % ILP == 0 &&
+       chunk_size % ILP == 0 &&
+       is_aligned(p) &&
+       is_aligned(update))
+    {
+      T r_p[ILP];
+      T r_update[ILP];
+      for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+      {
+        // load
+        load_store(r_p, p, 0, i_start);
+        load_store(r_update, update, 0, i_start);
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_p[ii] = static_cast<MATH_T>(r_p[ii]) - (ratio * static_cast<MATH_T>(r_update[ii]));
+        }
+        load_store(p, r_p, i_start, 0);
+      }
+    }
+    else
+    {
+      for(int i_start = 0;
+          i_start < n && i_start < chunk_size;
+          i_start += blockDim.x*ILP)
+      {
+        MATH_T r_p[ILP];
+        MATH_T r_update[ILP];
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            r_p[ii] = p[i];
+            r_update[ii] = update[i];
+          }
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_p[ii] = r_p[ii] - (ratio * r_update[ii]);
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+          {
+            p[i] = r_p[ii];
+          }
+        }
+      }
+    }
+  }
+};
+
+
+void multi_tensor_lamb_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  const float lr,
+  const float beta1,
+  const float beta2,
+  const float epsilon,
+  const int step,
+  const int bias_correction,
+  const float weight_decay,
+  const int grad_averaging,
+  const int mode,
+  at::Tensor global_grad_norm,
+  const float max_grad_norm,
+  at::optional<bool> use_nvlamb_python)
+{
+  using namespace at;
+  // Master weight and 32bit momentum(potentially changing) is not handled by this
+  // So we assume every tensor are all in the same type
+
+  bool use_nvlamb = use_nvlamb_python.has_value() ? use_nvlamb_python.value() : false;
+
+  // Handle bias correction mode
+  float bias_correction1 = 1.0f, bias_correction2 = 1.0f;
+  if (bias_correction == 1) {
+    bias_correction1 = 1 - std::pow(beta1, step);
+    bias_correction2 = 1 - std::pow(beta2, step);
+  }
+
+  // Handle grad averaging mode
+  float beta3 = 1.0f;
+  if (grad_averaging == 1) beta3 = 1 - beta1;
+
+  std::vector<std::vector<at::Tensor>> grad_list(tensor_lists.begin(), tensor_lists.begin()+1);
+  std::vector<std::vector<at::Tensor>> param_list(tensor_lists.begin()+1, tensor_lists.begin()+2);
+
+  // Compute per tensor param norm
+  auto param_norm_tuple = multi_tensor_l2norm_cuda(chunk_size, noop_flag, param_list, true);
+
+  // We now in-place modify grad to store update before compute its norm
+  // Generally this is not a issue since people modify grad in step() method all the time
+  // We can also grab list of empty tensor to avoid this, but I'd like to save space/cpu code
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "lamb_stage_1",
+      multi_tensor_apply<4>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        LAMBStage1Functor<scalar_t_0>(),
+        beta1,
+        beta2,
+        beta3, // 1-beta1 or 1 depends on averaging mode
+        bias_correction1,
+        bias_correction2,
+        epsilon,
+        (adamMode_t) mode,
+        weight_decay,
+        global_grad_norm.DATA_PTR<float>(),
+        max_grad_norm); )
+
+  // Compute update norms
+  auto update_norm_tuple = multi_tensor_l2norm_cuda(chunk_size, noop_flag, grad_list, true);
+
+  std::vector<std::vector<at::Tensor>> grad_param_list(tensor_lists.begin(), tensor_lists.begin()+2);
+
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "lamb_stage_2",
+      multi_tensor_apply<2>(
+        BLOCK_SIZE,
+        chunk_size,
+       	noop_flag,
+        grad_param_list,
+        LAMBStage2Functor<scalar_t_0>(),
+        std::get<1>(param_norm_tuple).DATA_PTR<float>(),
+        std::get<1>(update_norm_tuple).DATA_PTR<float>(),
+        lr,
+	weight_decay,
+	use_nvlamb); )
+
+  AT_CUDA_CHECK(cudaGetLastError());
+
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb_stage_1.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb_stage_1.cu
new file mode 100644
index 0000000000..6ad7649bc5
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb_stage_1.cu
@@ -0,0 +1,151 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+// Step 1 computes the 'update' value of regular Adam optimizer.
+template<typename GRAD_T, typename T, typename UPD_T>
+struct LAMBStage1Functor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<5>& tl,
+    const float* per_tensor_decay,
+    const float beta1,
+    const float beta2,
+    const float beta1_correction,
+    const float beta2_correction,
+    const float epsilon,
+    const float clipped_global_grad_norm)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int tensor_num = tl.start_tensor_this_launch + tensor_loc;
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    float decay = per_tensor_decay[tensor_num];
+
+    GRAD_T* g = (GRAD_T*)tl.addresses[0][tensor_loc];
+    g += chunk_idx*chunk_size;
+
+    T* p = (T*)tl.addresses[1][tensor_loc];
+    p += chunk_idx*chunk_size;
+
+    T* m = (T*)tl.addresses[2][tensor_loc];
+    m += chunk_idx*chunk_size;
+
+    T* v = (T*)tl.addresses[3][tensor_loc];
+    v += chunk_idx*chunk_size;
+
+    UPD_T* update = (UPD_T*)tl.addresses[4][tensor_loc];
+    update += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    // see note in multi_tensor_scale_kernel.cu
+    for(int i_start = 0;
+            i_start < n && i_start < chunk_size;
+            i_start += blockDim.x*ILP)
+    {
+      GRAD_T r_g[ILP];
+      T r_p[ILP];
+      T r_m[ILP];
+      T r_v[ILP];
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          r_g[ii] = g[i];
+          r_p[ii] = p[i];
+          r_m[ii] = m[i];
+          r_v[ii] = v[i];
+        } else {
+          r_g[ii] = GRAD_T(0);
+          r_p[ii] = T(0);
+          r_m[ii] = T(0);
+          r_v[ii] = T(0);
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        T scaled_grad = r_g[ii] / clipped_global_grad_norm;
+        r_m[ii] = r_m[ii] * beta1 + (1-beta1) * scaled_grad;
+        r_v[ii] = r_v[ii] * beta2 + (1-beta2) * scaled_grad * scaled_grad;
+        T next_m_unbiased = r_m[ii] / beta1_correction;
+        T next_v_unbiased = r_v[ii] / beta2_correction;
+        T denom = std::sqrt(next_v_unbiased) + epsilon;
+        r_p[ii] = (next_m_unbiased/denom) + (decay*r_p[ii]);
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          update[i] = (UPD_T)r_p[ii];
+          m[i] = r_m[ii];
+          v[i] = r_v[ii];
+        }
+      }
+    }
+  }
+};
+
+void multi_tensor_lamb_stage1_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor per_tensor_decay,
+  const int step,
+  const float beta1,
+  const float beta2,
+  const float epsilon,
+  at::Tensor global_grad_norm,
+  const float max_global_grad_norm)
+{
+  using namespace at;
+
+  const float* g_grad_norm = global_grad_norm.DATA_PTR<float>();
+  float clipped_global_grad_norm = *(g_grad_norm) > max_global_grad_norm ? *(g_grad_norm) / max_global_grad_norm : 1.0f;
+  float next_step = float(step+1);
+  float beta1_correction = 1.0f - std::pow(beta1, next_step);
+  float beta2_correction = 1.0f - std::pow(beta2, next_step);
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "lamb_stage_1",
+    DISPATCH_FLOAT_AND_HALF(tensor_lists[1][0].scalar_type(), 1, "lamb_stage_1",
+      DISPATCH_FLOAT_AND_HALF(tensor_lists[4][0].scalar_type(), 2, "lamb_stage_1",
+        multi_tensor_apply<5>(
+          BLOCK_SIZE,
+          chunk_size,
+          noop_flag,
+          tensor_lists,
+          LAMBStage1Functor<scalar_t_0, scalar_t_1, scalar_t_2>(),
+          per_tensor_decay.DATA_PTR<float>(),
+          beta1,
+          beta2,
+          beta1_correction,
+          beta2_correction,
+          epsilon,
+          clipped_global_grad_norm); )))
+
+  AT_CUDA_CHECK(cudaGetLastError());
+
+  // AT_CUDA_CHECK(cudaDeviceSynchronize());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb_stage_2.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb_stage_2.cu
new file mode 100644
index 0000000000..90970666c7
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_lamb_stage_2.cu
@@ -0,0 +1,125 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+using MATH_T = float;
+
+// Step 2 reads in 'update' value and per-tensor param_norm and update_norm.
+// It computes new parameter value.
+template<typename T, typename UPD_T>
+struct LAMBStage2Functor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<2>& tl,
+    const float* per_tensor_param_norm,
+    const float* per_tensor_update_norm,
+    const float learning_rate,
+    const float decay,
+    bool use_nvlamb)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int tensor_num = tl.start_tensor_this_launch + tensor_loc;
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    MATH_T ratio = learning_rate;
+    // nvlamb: apply adaptive learning rate to all parameters
+    // otherwise, only apply to those with non-zero weight decay
+    if (use_nvlamb || (decay != 0.0))
+    {
+      float param_norm = per_tensor_param_norm[tensor_num];
+      float update_norm = per_tensor_update_norm[tensor_num];
+      ratio = (update_norm != 0.0f && param_norm != 0.0f) ? learning_rate * (param_norm / update_norm) : learning_rate;
+    }
+
+    T* p = (T*)tl.addresses[0][tensor_loc];
+    p += chunk_idx*chunk_size;
+
+    UPD_T* update = (UPD_T*)tl.addresses[1][tensor_loc];
+    update += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    for(int i_start = 0;
+            i_start < n && i_start < chunk_size;
+            i_start += blockDim.x*ILP)
+    {
+      T r_p[ILP];
+      UPD_T r_update[ILP];
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          r_p[ii] = p[i];
+          r_update[ii] = update[i];
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        r_p[ii] = r_p[ii] - (ratio*(T)r_update[ii]);
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          p[i] = r_p[ii];
+        }
+      }
+    }
+  }
+};
+
+void multi_tensor_lamb_stage2_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor per_tensor_param_norm,
+  at::Tensor per_tensor_update_norm,
+  const float lr,
+  const float weight_decay,
+  at::optional<bool> use_nvlamb_python)
+{
+  bool use_nvlamb = use_nvlamb_python.has_value() ? use_nvlamb_python.value() : false;
+
+  using namespace at;
+
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "lamb_stage_2",
+    DISPATCH_FLOAT_AND_HALF(tensor_lists[1][0].scalar_type(), 1, "lamb_stage_2",
+      multi_tensor_apply<2>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        LAMBStage2Functor<scalar_t_0, scalar_t_1>(),
+        per_tensor_param_norm.DATA_PTR<float>(),
+        per_tensor_update_norm.DATA_PTR<float>(),
+        lr,
+	weight_decay,
+	use_nvlamb); ))
+
+  AT_CUDA_CHECK(cudaGetLastError());
+
+  // AT_CUDA_CHECK(cudaDeviceSynchronize());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_novograd.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_novograd.cu
new file mode 100644
index 0000000000..2decc06b8c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_novograd.cu
@@ -0,0 +1,188 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+typedef enum{
+  MOMENT_MODE_0   =0, // Novograd paper mode, momentum caculation with denom then decay inside
+  MOMENT_MODE_1   =1  // Decoupled weight decay mode
+} momentMode_t;
+
+void multi_tensor_norm_out_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor out,
+  const float alpha,
+  const float beta,
+  const int norm_type);
+
+using MATH_T = float;
+
+template<typename T>
+struct NovoGradFunctor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<3>& tl,
+    const float beta1,
+    const float beta2,
+    const float beta3,
+    const float beta1_correction,
+    const float beta2_correction,
+    const float epsilon,
+    const float lr,
+    momentMode_t m_mode,
+    const float decay,
+    const float* per_tensor_grad_norm)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int tensor_num = tl.start_tensor_this_launch + tensor_loc;
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    float grad_norm = per_tensor_grad_norm[tensor_num];
+
+    T* g = (T*)tl.addresses[0][tensor_loc];
+    g += chunk_idx*chunk_size;
+
+    T* p = (T*)tl.addresses[1][tensor_loc];
+    p += chunk_idx*chunk_size;
+
+    T* m = (T*)tl.addresses[2][tensor_loc];
+    m += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    // see note in multi_tensor_scale_kernel.cu
+    for(int i_start = 0;
+            i_start < n && i_start < chunk_size;
+            i_start += blockDim.x*ILP)
+    {
+      MATH_T r_g[ILP];
+      MATH_T r_p[ILP];
+      MATH_T r_m[ILP];
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          r_g[ii] = g[i];
+          r_p[ii] = p[i];
+          r_m[ii] = m[i];
+        } else {
+          r_g[ii] = MATH_T(0);
+          r_p[ii] = MATH_T(0);
+          r_m[ii] = MATH_T(0);
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        if (m_mode == MOMENT_MODE_0) {
+          MATH_T next_v_unbiased = grad_norm / beta2_correction;
+          MATH_T denom = next_v_unbiased + epsilon;
+          r_g[ii] = (r_g[ii] / denom) + (decay * r_p[ii]);
+          r_m[ii] = beta1 * r_m[ii] + beta3 * r_g[ii];
+          MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+          r_p[ii] = r_p[ii] - (lr * next_m_unbiased);
+        }
+        else {
+          r_m[ii] = beta1 * r_m[ii] + beta3 * r_g[ii];
+          MATH_T next_m_unbiased = r_m[ii] / beta1_correction;
+          MATH_T next_v_unbiased = grad_norm / beta2_correction;
+          MATH_T denom = next_v_unbiased + epsilon;
+          MATH_T update = (next_m_unbiased / denom) + (decay * r_p[ii]);
+          r_p[ii] = r_p[ii] - (lr * update);
+        }
+      }
+#pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          p[i] = r_p[ii];
+          m[i] = r_m[ii];
+        }
+      }
+    }
+  }
+};
+
+void multi_tensor_novograd_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  at::Tensor grad_norms,
+  const float lr,
+  const float beta1,
+  const float beta2,
+  const float epsilon,
+  const int step,
+  const int bias_correction,
+  const float weight_decay,
+  const int grad_averaging,
+  const int moment_mode,
+  const int norm_type)
+{
+  using namespace at;
+
+  // Handle bias correction mode
+  float bias_correction1 = 1.0f, bias_correction2 = 1.0f;
+  if (bias_correction == 1) {
+    bias_correction1 = 1 - std::pow(beta1, step);
+    bias_correction2 = std::sqrt(1 - std::pow(beta2, step));
+  }
+
+  // Handle grad averaging mode
+  float beta3 = 1;
+  if (grad_averaging == 1) beta3 = 1 - beta1;
+
+  std::vector<std::vector<at::Tensor>> grad_list(tensor_lists.begin(), tensor_lists.begin()+1);
+
+  // Compute and update grad norm
+  // Here use a per tensor norm, and blend new norm(n) and old norm(gn) by
+  // L-2: gn = sqrt(a * gn^2 + b * n^2)
+  // L-inf: gn = a * gn + b * n
+  multi_tensor_norm_out_cuda(chunk_size, noop_flag, grad_list, grad_norms, beta2, (1.0f - beta2), norm_type);
+
+  // Assume single type across p,g,m1,m2 now
+  DISPATCH_DOUBLE_FLOAT_AND_HALF(
+    tensor_lists[0][0].scalar_type(), 0, "novograd",
+    multi_tensor_apply<3>(
+      BLOCK_SIZE,
+      chunk_size,
+      noop_flag,
+      tensor_lists,
+      NovoGradFunctor<scalar_t_0>(),
+      beta1,
+      beta2,
+      beta3, // 1-beta1 or 1 depends on averaging mode
+      bias_correction1,
+      bias_correction2,
+      epsilon,
+      lr,
+      (momentMode_t) moment_mode,
+      weight_decay,
+      grad_norms.DATA_PTR<float>()); )
+
+  AT_CUDA_CHECK(cudaGetLastError());
+
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_scale_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_scale_kernel.cu
new file mode 100644
index 0000000000..629ee94207
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_scale_kernel.cu
@@ -0,0 +1,136 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+// Another possibility:
+// #include <torch/all.h>
+
+#include <assert.h>
+// Stringstream is a big hammer, but I want to rely on operator<< for dtype.
+#include <sstream>
+
+#include "type_shim.h"
+#include "multi_tensor_apply.cuh"
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+template<typename T>
+__device__ __forceinline__ bool is_aligned(T* p){
+  return ((uint64_t)p) % (ILP*sizeof(T)) == 0;
+}
+
+template<typename T>
+__device__ __forceinline__ void load_store(T* dst, T* src, int dst_offset, int src_offset){
+  typedef typename std::aligned_storage<ILP*sizeof(T), ILP*alignof(T)>::type LT;
+  ((LT*)dst)[dst_offset] = ((LT*)src)[src_offset];
+}
+
+template<typename in_t, typename out_t>
+struct ScaleFunctor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<2>& tl,
+    float scale)
+  {
+    // I'd like this kernel to propagate infs/nans.
+    // if(*noop_gmem == 1)
+    //   return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    in_t* in = (in_t*)tl.addresses[0][tensor_loc];
+    in += chunk_idx*chunk_size;
+
+    out_t* out = (out_t*)tl.addresses[1][tensor_loc];
+    out += chunk_idx*chunk_size;
+
+    n -= chunk_idx*chunk_size;
+
+    bool finite = true;
+    in_t r_in[ILP];
+    out_t r_out[ILP];
+
+    // to make things simple, we put aligned case in a different code path
+    if(n % ILP == 0 && chunk_size % ILP == 0 && is_aligned(in) && is_aligned(out))
+    {
+      for(int i_start = threadIdx.x; i_start*ILP < n && i_start*ILP < chunk_size; i_start += blockDim.x)
+      {
+        // load
+        load_store(r_in, in, 0 , i_start);
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_out[ii] = static_cast<float>(r_in[ii]) * scale;
+          finite = finite && isfinite(r_in[ii]);
+        }
+        // store
+        load_store(out, r_out, i_start, 0);
+      }
+    }
+    else
+    {
+      // Non-divergent exit condition for __syncthreads, not necessary here
+      for(int i_start = 0; i_start < n && i_start < chunk_size; i_start += blockDim.x*ILP)
+      {
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_in[ii] = 0;
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+            r_in[ii] = in[i];
+        }
+        // note for clarification to future michael:
+        // From a pure memory dependency perspective, there's likely no point unrolling
+        // the write loop, since writes just fire off once their LDGs arrive.
+        // Put another way, the STGs are dependent on the LDGs, but not on each other.
+        // There is still compute ILP benefit from unrolling the loop though.
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          r_out[ii] = static_cast<float>(r_in[ii]) * scale;
+          finite = finite && isfinite(r_in[ii]);
+        }
+#pragma unroll
+        for(int ii = 0; ii < ILP; ii++)
+        {
+          int i = i_start + threadIdx.x + ii*blockDim.x;
+          if(i < n && i < chunk_size)
+            out[i] = r_out[ii];
+        }
+      }
+    }
+    if(!finite)
+      *noop_gmem = 1; // Blindly fire off a write.  These will race but that's ok.
+  }
+};
+
+void multi_tensor_scale_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  float scale)
+{
+  using namespace at;
+  // The output (downscaled) type is always float.
+  // If build times suffer, think about where to put this dispatch,
+  // and what logic should be moved out of multi_tensor_apply.
+
+  DISPATCH_FLOAT_AND_HALF(tensor_lists[0][0].scalar_type(), 0, "multi_tensor_scale_cuda",
+    DISPATCH_FLOAT_AND_HALF(tensor_lists[1][0].scalar_type(), 1, "multi_tensor_scale_cuda",
+      multi_tensor_apply<2>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        ScaleFunctor<scalar_t_0, scalar_t_1>(),
+        scale); ))
+  AT_CUDA_CHECK(cudaGetLastError());
+
+  // AT_CUDA_CHECK(cudaDeviceSynchronize());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_sgd_kernel.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_sgd_kernel.cu
new file mode 100644
index 0000000000..42a7406be6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/multi_tensor_sgd_kernel.cu
@@ -0,0 +1,280 @@
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+#include <ATen/cuda/Exceptions.h>
+#include "multi_tensor_apply.cuh"
+#include "compat.h"
+
+#include <assert.h>
+#include <cuda_runtime.h>
+
+#define BLOCK_SIZE 512
+#define ILP 4
+
+/**
+ * Perform fused SGD on multiple buffers
+ * N: number of tensors
+ * tl[0] : gradients
+ * tl[1] : weights
+ * tl[2] : momentum buffers
+ * tl[3] : fp16 weights (if appropriate)
+ * wd : weight_decay (scalar)
+ * momentum : momentum (scalar)
+ * dampening : momentum dampening (scalar)
+ * lr : learning rate (scalar)
+ * nesterov : enable nesterov (bool)
+ * first run : necessary for proper momentum handling & init
+ * wd_after_momentum : apply weight decay _after_ momentum instead of before
+ **/
+template<int N, typename T_grad, typename T_weight>
+struct SGDFunctor
+{
+   __device__ __forceinline__ void operator()(
+    int chunk_size,
+    volatile int* noop_gmem,
+    TensorListMetadata<N>& tl,
+    float wd,
+    float momentum,
+    float dampening,
+    float lr,
+    bool nesterov,
+    bool first_run,
+    bool wd_after_momentum,
+    float scale)
+  {
+    // Early exit if we don't need to do anything
+    if (*noop_gmem) return;
+
+    int tensor_loc = tl.block_to_tensor[blockIdx.x];
+    int chunk_idx = tl.block_to_chunk[blockIdx.x];
+    int n = tl.sizes[tensor_loc];
+
+    T_grad* grad_in = (T_grad*)tl.addresses[0][tensor_loc];
+    grad_in += chunk_idx*chunk_size;
+
+    T_weight* weight_in = (T_weight*)tl.addresses[1][tensor_loc];
+    weight_in += chunk_idx*chunk_size;
+
+    T_weight* mom_in = (T_weight*)tl.addresses[2][tensor_loc];
+    mom_in += chunk_idx*chunk_size;
+
+    at::Half *model_weights_out = nullptr;
+    if(N == 4)
+    {
+      model_weights_out = (at::Half*)tl.addresses[3][tensor_loc];
+      model_weights_out += chunk_idx*chunk_size;
+    }
+
+    n -= chunk_idx*chunk_size;
+
+    // Non-divergent exit condition for the __syncthreads
+    float incoming_grads[ILP];
+    float incoming_weights[ILP];
+    float incoming_moms[ILP];
+    for(int i_start = 0;
+        i_start < n && i_start < chunk_size;
+        i_start += blockDim.x*ILP)
+    {
+      #pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        incoming_grads[ii] = 0;
+        incoming_weights[ii] = 0;
+        incoming_moms[ii] = 0;
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          incoming_grads[ii] = static_cast<float>(grad_in[i])*scale;
+          incoming_weights[ii] = static_cast<float>(weight_in[i]);
+          incoming_moms[ii] = static_cast<float>(mom_in[i]);
+        }
+      }
+
+      // note for clarification to future michael:
+      // From a pure memory dependency perspective, there's likely no point unrolling
+      // the write loop, since writes just fire off once their LDGs arrive.
+      // Put another way, the STGs are dependent on the LDGs, but not on each other.
+      // There is still compute ILP benefit from unrolling the loop though.
+      #pragma unroll
+      for(int ii = 0; ii < ILP; ii++)
+      {
+        int i = i_start + threadIdx.x + ii*blockDim.x;
+        if(i < n && i < chunk_size)
+        {
+          // apply weight decay before momentum if necessary
+          if(wd != 0.f && !wd_after_momentum)
+            incoming_grads[ii] += wd * incoming_weights[ii];
+
+          if(momentum != 0.f)
+          {
+            if(!first_run)
+              incoming_moms[ii] = incoming_moms[ii] * momentum + (1.f - dampening) * incoming_grads[ii];
+            else // initialize momentums to current incoming grads
+              incoming_moms[ii] = incoming_grads[ii];
+
+            if(nesterov)
+              incoming_grads[ii] += momentum * incoming_moms[ii];
+            else
+              incoming_grads[ii] = incoming_moms[ii];
+          }
+
+          // Apply WD after momentum if desired
+          if(wd != 0.f && wd_after_momentum)
+            incoming_grads[ii] += wd * incoming_weights[ii];
+
+          // adjust the weight and write out
+          weight_in[i] += (-lr * incoming_grads[ii]);
+
+          // if necessary, write out an fp16 copy of the weights
+          if(N == 4)
+            model_weights_out[i] = static_cast<at::Half>(weight_in[i]);
+
+          // also write out the new momentum
+          if(momentum != 0.f)
+            mom_in[i] = incoming_moms[ii];
+        }
+      }
+    }
+  }
+};
+
+void multi_tensor_sgd_cuda(
+  int chunk_size,
+  at::Tensor noop_flag,
+  std::vector<std::vector<at::Tensor>> tensor_lists,
+  float wd,
+  float momentum,
+  float dampening,
+  float lr,
+  bool nesterov,
+  bool first_run,
+  bool wd_after_momentum,
+  float scale)
+{
+  auto num_tensors = tensor_lists.size();
+  auto grad_type = tensor_lists[0][0].scalar_type();
+  auto weight_type = tensor_lists[1][0].scalar_type();
+
+  if(num_tensors == 4)
+    for(int i = 0; i < tensor_lists[3].size(); i++)
+        TORCH_CHECK(tensor_lists[3][i].scalar_type() == at::ScalarType::Half,
+                 "Additional output tensors should always be fp16.");
+
+  TORCH_CHECK(noop_flag.device() == tensor_lists[0][0].device(), "expected noop flag to be on the same device as tensors");
+
+  // We have 3 possibilities to handle here, in terms of
+  // grad_type, param_type, momentum_type, requires_fp16_copy
+  // 1. fp16, fp16, fp16, No
+  // 2. fp32, fp32, fp32, No
+  // 3. fp16, fp32, fp32, Yes
+  // 4. fp32, fp32, fp32, Yes // this is the materialize_master_grads=True case
+  // It's easier to hardcode these possibilities than to use
+  // switches etc. to handle the cross-product of cases where
+  // we don't want the majority of them.
+
+  // Case 1. fp16, fp16, fp16, No
+  if(grad_type == at::ScalarType::Half &&
+     weight_type == at::ScalarType::Half &&
+     num_tensors == 3)
+  {
+    multi_tensor_apply<3>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        SGDFunctor<3, at::Half, at::Half>(),
+        wd,
+        momentum,
+        dampening,
+        lr,
+        nesterov,
+        first_run,
+        wd_after_momentum,
+        scale);
+  }
+  // Case 2. fp16, fp32, fp32, No
+  // else if (grad_type == at::ScalarType::Half &&
+  //          weight_type == at::ScalarType::Float &&
+  //          num_tensors == 3) {
+  //   multi_tensor_apply<3>(
+  //       BLOCK_SIZE,
+  //       chunk_size,
+  //       noop_flag,
+  //       tensor_lists,
+  //       SGDFunctor<3, at::Half, float>(),
+  //       wd,
+  //       momentum,
+  //       dampening,
+  //       lr,
+  //       nesterov,
+  //       first_run,
+  //       wd_after_momentum);
+  // }
+  // Case 2. fp32, fp32, fp32, No
+  else if(grad_type == at::ScalarType::Float &&
+          weight_type == at::ScalarType::Float &&
+          num_tensors == 3)
+  {
+    multi_tensor_apply<3>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        SGDFunctor<3, float, float>(),
+        wd,
+        momentum,
+        dampening,
+        lr,
+        nesterov,
+        first_run,
+        wd_after_momentum,
+        scale);
+  }
+  // Case 3. fp16, fp32, fp32, Yes
+  else if(grad_type == at::ScalarType::Half &&
+          weight_type == at::ScalarType::Float &&
+          num_tensors == 4)
+  {
+    multi_tensor_apply<4>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        SGDFunctor<4, at::Half, float>(),
+        wd,
+        momentum,
+        dampening,
+        lr,
+        nesterov,
+        first_run,
+        wd_after_momentum,
+        scale);
+  }
+  // Case 4. fp32, fp32, fp32, Yes
+  else if(grad_type == at::ScalarType::Float &&
+          weight_type == at::ScalarType::Float &&
+          num_tensors == 4)
+  {
+    multi_tensor_apply<4>(
+        BLOCK_SIZE,
+        chunk_size,
+        noop_flag,
+        tensor_lists,
+        SGDFunctor<4, float, float>(),
+        wd,
+        momentum,
+        dampening,
+        lr,
+        nesterov,
+        first_run,
+        wd_after_momentum,
+        scale);
+  }
+  else
+  {
+    AT_ERROR("multi_tensor_sgd only supports some combinations of gradient & weight types. Given: ",
+             "gradient: ", grad_type, ", weight: ", weight_type, ", num_lists: ", num_tensors);
+  }
+
+  AT_CUDA_CHECK(cudaGetLastError());
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/syncbn.cpp b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/syncbn.cpp
new file mode 100644
index 0000000000..578a6e653b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/syncbn.cpp
@@ -0,0 +1,109 @@
+#include <torch/extension.h>
+#include <ATen/ATen.h>
+
+#include <vector>
+
+// returns {mean,biased_var}
+// implemented using welford 
+std::vector<at::Tensor> welford_mean_var_CUDA(const at::Tensor input);
+
+// reduces array of mean/var across processes
+// returns global {mean,inv_std,biased_var}
+// implemented using welford 
+std::vector<at::Tensor> welford_parallel_CUDA(const at::Tensor mean_feature_nodes,
+                                              const at::Tensor var_biased_feature_nodes,
+                                              const at::Tensor numel,
+                                              const float eps);
+
+// elementwise BN operation, returns output
+// input/weight/shift should have identical data type;
+// mean/inv_std have promoted data type (dtype==fp16?fp32:dtype)
+at::Tensor batchnorm_forward_CUDA(const at::Tensor input,
+                                  const at::Tensor mean,
+                                  const at::Tensor inv_std,
+                                  const at::optional<at::Tensor> weight,
+                                  const at::optional<at::Tensor> shift);
+
+// backward BN operation, returns {sum_dy, sum_dy_xmu, grad_weight, grad_bias}
+// grad_output/input should have identical data type;
+// mean/inv_std have promoted data type (dtype==fp16?fp32:dtype)
+// implemented using kahan summation
+std::vector<at::Tensor> reduce_bn_CUDA(const at::Tensor grad_output,
+                                           const at::Tensor input,
+                                           const at::Tensor mean,
+                                           const at::Tensor inv_std,
+                                           const at::optional<at::Tensor> weight);
+
+// elementwise backward BN operation, returns grad_input
+// grad_output/input/weight precision could be fp16/fp32;
+// mean/inv_std/sum_dy/sum_dy_xmu precision is fp32
+at::Tensor batchnorm_backward_CUDA(const at::Tensor grad_output,
+                                   const at::Tensor input,
+                                   const at::Tensor mean,
+                                   const at::Tensor inv_std,
+                                   const at::optional<at::Tensor> weight,
+                                   const at::Tensor sum_dy,
+                                   const at::Tensor sum_dy_xmu,
+                                   const at::Tensor count);
+
+// returns {mean, biased_var}
+// implemented using welford 
+// expect data to be in n+c format (channel last) and applies CUDNN_BATCHNORM_SPATIAL
+std::vector<at::Tensor> welford_mean_var_c_last_CUDA(const at::Tensor input);
+
+// elementwise BN operation, returns output
+// input/weight/shift should have identical data type;
+// mean/inv_std have promoted data type (dtype==fp16?fp32:dtype)
+// expect data to be in n+c format (channel last) and applies CUDNN_BATCHNORM_SPATIAL
+at::Tensor batchnorm_forward_c_last_CUDA(const at::Tensor input,
+                                         const at::optional<at::Tensor> z,
+                                         const at::Tensor mean,
+                                         const at::Tensor inv_std,
+                                         const at::optional<at::Tensor> weight,
+                                         const at::optional<at::Tensor> shift,
+                                         const bool fuse_relu);
+
+// backward BN operation, returns {sum_dy, sum_dy_xmu, grad_weight, grad_bias}
+// grad_output/input should have identical data type;
+// mean/inv_std have promoted data type (dtype==fp16?fp32:dtype)
+// expect data to be in n+c format (channel last) and applies CUDNN_BATCHNORM_SPATIAL
+std::vector<at::Tensor> reduce_bn_c_last_CUDA(const at::Tensor grad_output,
+                                              const at::Tensor input,
+                                              const at::Tensor mean,
+                                              const at::Tensor inv_std,
+                                              const at::optional<at::Tensor> weight);
+
+// elementwise backward BN operation, returns grad_input
+// grad_output/input/weight precision could be fp16/fp32;
+// mean/inv_std/sum_dy/sum_dy_xmu precision is fp32
+// expect data to be in n+c format (channel last) and applies CUDNN_BATCHNORM_SPATIAL
+at::Tensor batchnorm_backward_c_last_CUDA(const at::Tensor grad_output,
+                                          const at::Tensor input,
+                                          const at::Tensor mean,
+                                          const at::Tensor inv_std,
+                                          const at::optional<at::Tensor> weight,
+                                          const at::Tensor sum_dy,
+                                          const at::Tensor sum_dy_xmu,
+                                          const at::Tensor count);
+
+at::Tensor relu_backward_c_last_CUDA(const at::Tensor grad_output,
+                                     const at::Tensor input,
+                                     const at::optional<at::Tensor> z,
+                                     const at::Tensor mean,
+                                     const at::Tensor inv_std,
+                                     const at::optional<at::Tensor> weight,
+                                     const at::optional<at::Tensor> shift);
+
+
+PYBIND11_MODULE(TORCH_EXTENSION_NAME, m) {
+  m.def("welford_mean_var", &welford_mean_var_CUDA, "welford mean variance");
+  m.def("welford_parallel", &welford_parallel_CUDA, "welford parallel reduce mean variance");
+  m.def("batchnorm_forward", &batchnorm_forward_CUDA, "batchnorm forward");
+  m.def("reduce_bn", &reduce_bn_CUDA, "batchnorm backward reduce grad sum and bias/weight grad");
+  m.def("batchnorm_backward", &batchnorm_backward_CUDA, "batchnorm backward dgrad");
+  m.def("welford_mean_var_c_last", &welford_mean_var_c_last_CUDA, "welford mean variance nhwc");
+  m.def("batchnorm_forward_c_last", &batchnorm_forward_c_last_CUDA, "batchnorm forward nhwc");
+  m.def("reduce_bn_c_last", &reduce_bn_c_last_CUDA, "batchnorm backwards reduce grad sum and bias/weight grad nhwc");
+  m.def("batchnorm_backward_c_last", &batchnorm_backward_c_last_CUDA, "batchnorm backward dgrad nhwc");
+  m.def("relu_bw_c_last", &relu_backward_c_last_CUDA, "relu_bw_c_last");
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/type_shim.h b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/type_shim.h
new file mode 100644
index 0000000000..b9d03a227b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/type_shim.h
@@ -0,0 +1,387 @@
+#include <ATen/ATen.h>
+#include "compat.h"
+
+// Forward/backward compatiblity hack around
+// https://github.com/pytorch/pytorch/commit/3aeb78079bcd68282fe9117088e138b77318e288
+// pending more future-proof guidance from upstream.
+// struct TypeShim
+// {
+//   const at::Type& payload;
+//   TypeShim(const at::Type& type) : payload(type) {}
+//   // Enable trivial conversion to a const at::Type& for pre-3aeb78
+//   operator const at::Type&(){ return payload; };
+//   // Enable dispatch switch statements to take *this directly for  post-3aeb78
+//   //operator at::ScalarType(){ return payload.; };
+// };
+
+#define DISPATCH_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \
+  switch(TYPE) \
+  { \
+    case at::ScalarType::Float: \
+    { \
+      using scalar_t_##LEVEL = float; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    case at::ScalarType::Half: \
+    { \
+      using scalar_t_##LEVEL = at::Half; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    default: \
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'");  \
+  }
+
+
+#define DISPATCH_FLOAT_HALF_AND_BFLOAT(TYPE, LEVEL, NAME, ...) \
+  switch(TYPE) \
+  { \
+    case at::ScalarType::Float: \
+    { \
+      using scalar_t_##LEVEL = float; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    case at::ScalarType::Half: \
+    { \
+      using scalar_t_##LEVEL = at::Half; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    case at::ScalarType::BFloat16: \
+    { \
+      using scalar_t_##LEVEL = at::BFloat16; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    default: \
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'");  \
+  }
+
+
+#define DISPATCH_FLOAT_HALF_AND_BYTE(TYPE, LEVEL, NAME, ...) \
+  switch(TYPE) \
+  { \
+    case at::ScalarType::Float: \
+    { \
+      using scalar_t_##LEVEL = float; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    case at::ScalarType::Half: \
+    { \
+      using scalar_t_##LEVEL = at::Half; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    case at::ScalarType::Byte: \
+    { \
+      using scalar_t_##LEVEL = uint8_t; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    default: \
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'");  \
+  }
+
+
+#define DISPATCH_DOUBLE_FLOAT_AND_HALF(TYPE, LEVEL, NAME, ...) \
+  switch(TYPE) \
+  { \
+    case at::ScalarType::Double: \
+    { \
+      using scalar_t_##LEVEL = double; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    case at::ScalarType::Float: \
+    { \
+      using scalar_t_##LEVEL = float; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    case at::ScalarType::Half: \
+    { \
+      using scalar_t_##LEVEL = at::Half; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    default: \
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'");  \
+  }
+
+
+  #define DISPATCH_DOUBLE_AND_FLOAT(TYPE, LEVEL, NAME, ...) \
+  switch(TYPE) \
+  { \
+    case at::ScalarType::Double: \
+    { \
+      using scalar_t_##LEVEL = double; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    case at::ScalarType::Float: \
+    { \
+      using scalar_t_##LEVEL = float; \
+      __VA_ARGS__; \
+      break; \
+    } \
+    default: \
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'");  \
+  }
+
+
+  #define DISPATCH_HALF_AND_BFLOAT(TYPE, NAME, ...)			\
+  switch(TYPE)								\
+    {									\
+    case at::ScalarType::Half:						\
+      {									\
+	using scalar_t = at::Half;					\
+	__VA_ARGS__;							\
+	break;								\
+      }									\
+    case at::ScalarType::BFloat16:					\
+      {									\
+	using scalar_t = at::BFloat16;					\
+	__VA_ARGS__;							\
+	break;								\
+      }									\
+    default:								\
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPE), "'");	\
+  }
+
+
+  #define DISPATCH_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(TYPEIN, TYPEOUT, NAME, ...) \
+  switch(TYPEIN)							\
+    {									\
+    case at::ScalarType::Float:						\
+      {									\
+	using scalar_t_in = float;					\
+	switch(TYPEOUT)							\
+	  {								\
+	  case at::ScalarType::Float:					\
+	    {								\
+	      using scalar_t_out = float;				\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  case at::ScalarType::Half:					\
+	    {								\
+	      using scalar_t_out = at::Half;				\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  case at::ScalarType::BFloat16:				\
+	    {								\
+	      using scalar_t_out = at::BFloat16;			\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  default:							\
+	    AT_ERROR(#NAME, " not implemented for '", toString(TYPEOUT), "'"); \
+	  }								\
+	break;								\
+      }									\
+    case at::ScalarType::Half:						\
+      {									\
+	using scalar_t_in = at::Half;					\
+	using scalar_t_out = at::Half;					\
+	__VA_ARGS__;							\
+	break;								\
+      }									\
+    case at::ScalarType::BFloat16:					\
+      {									\
+	using scalar_t_in = at::BFloat16;				\
+	using scalar_t_out = at::BFloat16;				\
+	__VA_ARGS__;							\
+	break;								\
+      }									\
+    default:								\
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPEIN), "'");	\
+    }
+
+
+  #define DISPATCH_DOUBLE_FLOAT_HALF_AND_BFLOAT_INOUT_TYPES(TYPEIN, TYPEOUT, NAME, ...) \
+  switch(TYPEIN)							\
+    {									\
+    case at::ScalarType::Double:						\
+      {									\
+	using scalar_t_in = double;					\
+	switch(TYPEOUT)							\
+	  {								\
+	  case at::ScalarType::Double:					\
+	    {								\
+	      using scalar_t_out = double;				\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  case at::ScalarType::Float:					\
+	    {								\
+	      using scalar_t_out = float;				\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  case at::ScalarType::Half:					\
+	    {								\
+	      using scalar_t_out = at::Half;				\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  case at::ScalarType::BFloat16:				\
+	    {								\
+	      using scalar_t_out = at::BFloat16;			\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  default:							\
+	    AT_ERROR(#NAME, " not implemented for '", toString(TYPEOUT), "'"); \
+	  }								\
+	break;								\
+      }									\
+    case at::ScalarType::Float:						\
+      {									\
+	using scalar_t_in = float;					\
+	switch(TYPEOUT)							\
+	  {								\
+	  case at::ScalarType::Float:					\
+	    {								\
+	      using scalar_t_out = float;				\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  case at::ScalarType::Half:					\
+	    {								\
+	      using scalar_t_out = at::Half;				\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  case at::ScalarType::BFloat16:				\
+	    {								\
+	      using scalar_t_out = at::BFloat16;			\
+	      __VA_ARGS__;						\
+	      break;							\
+	    }								\
+	  default:							\
+	    AT_ERROR(#NAME, " not implemented for '", toString(TYPEOUT), "'"); \
+	  }								\
+	break;								\
+      }									\
+    case at::ScalarType::Half:						\
+      {									\
+	using scalar_t_in = at::Half;					\
+	using scalar_t_out = at::Half;					\
+	__VA_ARGS__;							\
+	break;								\
+      }									\
+    case at::ScalarType::BFloat16:					\
+      {									\
+	using scalar_t_in = at::BFloat16;				\
+	using scalar_t_out = at::BFloat16;				\
+	__VA_ARGS__;							\
+	break;								\
+      }									\
+    default:								\
+      AT_ERROR(#NAME, " not implemented for '", toString(TYPEIN), "'");	\
+    }
+
+
+template<typename T>
+__device__ __forceinline__ T reduce_block_into_lanes
+  (T *x,
+   T val,
+   int lanes=1,
+   bool share_result=false) // lanes is intended to be <= 32.
+{
+  int tid = threadIdx.x + threadIdx.y*blockDim.x;
+  int blockSize = blockDim.x*blockDim.y; // blockSize is intended to be a multiple of 32.
+
+  if(blockSize >= 64)
+  {
+    x[tid] = val;
+    __syncthreads();
+  }
+
+  #pragma unroll
+  for(int i = (blockSize >> 1); i >= 64; i >>= 1)
+  {
+    if(tid < i)
+      x[tid] = x[tid] + x[tid+i];
+    __syncthreads();
+  }
+
+  T final;
+
+  if(tid < 32)
+  {
+    if(blockSize >= 64)
+      final = x[tid] + x[tid+32];
+    else
+      final = val;
+    // __SYNCWARP();
+
+    #pragma unroll
+    for(int i = 16; i >= lanes; i >>= 1)
+      final = final + __shfl_down_sync(0xffffffff, final, i);
+  }
+
+  if(share_result)
+  {
+    if(tid < lanes)
+      x[tid] = final; // EpilogueOp
+    // Make sure the smem result is visible to all warps.
+    __syncthreads();
+  }
+
+  return final;
+}
+
+template<typename T>
+__device__ __forceinline__ T reduce_block_into_lanes_max_op
+  (T *x,
+   T val,
+   int lanes=1,
+   bool share_result=false) // lanes is intended to be <= 32.
+{
+  int tid = threadIdx.x + threadIdx.y*blockDim.x;
+  int blockSize = blockDim.x*blockDim.y; // blockSize is intended to be a multiple of 32.
+
+  if(blockSize >= 64)
+  {
+    x[tid] = val;
+    __syncthreads();
+  }
+
+  #pragma unroll
+  for(int i = (blockSize >> 1); i >= 64; i >>= 1)
+  {
+    if(tid < i)
+      x[tid] = fmaxf(fabsf(x[tid]), fabsf(x[tid+i]));
+    __syncthreads();
+  }
+
+  T final;
+
+  if(tid < 32)
+  {
+    if(blockSize >= 64)
+      final = fmaxf(fabsf(x[tid]), fabsf(x[tid+32]));
+    else
+      final = val;
+    // __SYNCWARP();
+
+    #pragma unroll
+    for(int i = 16; i >= lanes; i >>= 1)
+      final = fmaxf(fabsf(final), fabsf(__shfl_down_sync(0xffffffff, final, i)));
+  }
+
+  if(share_result)
+  {
+    if(tid < lanes)
+      x[tid] = final; // EpilogueOp
+    // Make sure the smem result is visible to all warps.
+    __syncthreads();
+  }
+
+  return final;
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/welford.cu b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/welford.cu
new file mode 100644
index 0000000000..374a3845ba
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/csrc/welford.cu
@@ -0,0 +1,1510 @@
+#include <iostream>
+#include <ATen/ATen.h>
+#include <ATen/AccumulateType.h>
+#include <ATen/cuda/CUDAContext.h>
+
+#include <cuda.h>
+#include <cuda_runtime.h>
+
+#include <vector>
+
+#include "type_shim.h"
+#include "compat.h"
+
+
+__device__ __forceinline__ int lastpow2(int n)
+{
+  int out = 1 << (31 - __clz(n));
+  if(n == out)
+    out >>= 1;
+  return out;
+}
+
+__host__ __forceinline__ int h_next_pow2(unsigned int n) {
+    n--;
+    n |= (n >>  1);
+    n |= (n >>  2);
+    n |= (n >>  4);
+    n |= (n >>  8);
+    n |= (n >> 16);
+    return ++n;
+}
+
+__host__ __forceinline__ int h_last_pow2(unsigned int n) {
+    n |= (n >>  1);
+    n |= (n >>  2);
+    n |= (n >>  4);
+    n |= (n >>  8);
+    n |= (n >> 16);
+    return n - (n >> 1);
+}
+
+
+#define WARP_SIZE 32
+
+template<typename T>
+__device__ __forceinline__ T warp_reduce_sum(T val)
+{
+  #pragma unroll
+  for(int i = WARP_SIZE/2; i > 0; i >>= 1)
+    val = val + __shfl_down_sync(0xffffffff, val, i);
+  return val;
+}
+
+template<typename T>
+__device__ __forceinline__ T reduce_block(T *x, T val)
+{
+  int tid = threadIdx.y*blockDim.x + threadIdx.x;
+  int blockSize = blockDim.x * blockDim.y;
+
+  if (blockSize > 32) {
+    val = warp_reduce_sum(val);
+    if (tid % WARP_SIZE == 0)
+      x[tid/WARP_SIZE] = val;
+
+    __syncthreads();
+
+    val = (tid < blockSize / WARP_SIZE? x[tid%WARP_SIZE] : T(0));
+  }
+
+  if(tid/WARP_SIZE==0) val = warp_reduce_sum(val);
+
+  return val;
+}
+
+#define ELEMENTS_PER_ITER 4 // enables concurrency within each thread to hide latency
+#define ELEMENTS_PER_THREAD 16
+#define OPTIMAL_TILE_W 32
+#define MAX_H_BLOCK 128
+#define MAX_BLOCK_SIZE 512
+
+__host__ int div_ru(int x, int y) {
+  return h_last_pow2(1 + (x-1)/y);
+}
+
+__host__ void flexible_launch_configs(
+      const int reduction,
+      const int stride,
+      dim3 &block,
+      dim3 &grid,
+      const bool coop_flag = false) {
+  int block_x = std::min(h_last_pow2(stride), OPTIMAL_TILE_W);
+  int block_y = std::min(h_last_pow2(div_ru(reduction , ELEMENTS_PER_THREAD)),
+                         MAX_BLOCK_SIZE / block_x);
+  if (block_x * block_y != MAX_BLOCK_SIZE) {
+    block_x = std::min(h_last_pow2(stride), MAX_BLOCK_SIZE / block_y);
+  }
+
+  int grid_x = div_ru(stride, block_x);
+  int grid_y = std::min(div_ru(reduction, block_y * ELEMENTS_PER_THREAD), MAX_H_BLOCK);
+  if (coop_flag) {
+    // it's not worth having a grid reduction if the reduction dimension is not big enough
+    grid_y = grid_y < 8 ? 1 : grid_y;
+  }
+
+  block.x = block_x;
+  block.y = block_y;
+  block.z = 1;
+  grid.x = grid_x;
+  grid.y = grid_y;
+  grid.z = 1;
+}
+
+template<typename T, typename C>
+__device__ __forceinline__ void welford_merge_element(C& count,
+                                                      T& mean,
+                                                      T& m2n,
+                                                      const C& num_new,
+                                                      const T& mean_new,
+                                                      const T& m2n_new) {
+      T factor = T(1.0) / max(1, (count + num_new));
+      T delta0 = mean - mean_new;
+      mean = (mean_new * num_new + mean * count) * factor;
+      m2n += m2n_new + delta0 * delta0 * num_new * count * factor;
+      count += num_new;
+}
+
+template<typename T>
+__device__ __forceinline__ void warp_reduce_mean_m2n(T &mean, T &m2n, int &num)
+{
+  #pragma unroll
+  for(int i = WARP_SIZE/2; i > 0; i >>= 1) {
+    auto num_new = __shfl_down_sync(0xffffffff, num, i);
+    auto mean_new = __shfl_down_sync(0xffffffff, mean, i);
+    auto m2n_new = __shfl_down_sync(0xffffffff, m2n, i);
+    welford_merge_element(num, mean, m2n, num_new, mean_new, m2n_new);
+  }
+}
+
+template <typename T>
+__device__ void welford_reduce_mean_m2n(
+      T* __restrict__ x,
+      int* __restrict__ count,
+      T &mean,
+      T &m2n,
+      int &num,
+      int block_size,
+      int thread_id)
+{
+  int lane = thread_id % WARP_SIZE;
+  int wid = thread_id / WARP_SIZE;
+
+  if (block_size > 32) {
+    warp_reduce_mean_m2n(mean, m2n, num);
+    if (lane == 0) {
+      x[wid*2] = mean;
+      x[wid*2+1] = m2n;
+      count[wid] = num;
+    }
+    __syncthreads();
+
+    if (wid == 0) {
+      mean = (thread_id < block_size / WARP_SIZE)? x[lane*2] : T(0);
+      m2n = (thread_id < block_size / WARP_SIZE)? x[lane*2+1] : T(0);
+      num = (thread_id < block_size / WARP_SIZE)? count[lane] : int(0);
+    }
+  }
+
+  if (wid==0) warp_reduce_mean_m2n(mean, m2n, num);
+
+  return;
+}
+
+// return spatial size for NC+ Tensors
+__host__ int get_tensor_spatial_size(const at::Tensor& input)
+{
+  auto space_size = input.size(2);
+  for (int i = 3; i < input.ndimension(); i++) {
+    space_size *= input.size(i);
+  }
+  return space_size;
+}
+
+// promote accumulation scalar type. promote half to float.
+__host__ at::ScalarType promote_scalartype(const at::Tensor& input)
+{
+  return input.scalar_type() == at::ScalarType::Half ?
+           at::ScalarType::Float : input.scalar_type();
+}
+
+// return single element size, optional accumulation type promotion.
+__host__ size_t get_element_data_size(const at::Tensor& input, bool accumulation = false)
+{
+  auto scalar_type = accumulation ? promote_scalartype(input) : input.scalar_type();
+  return at::elementSize(scalar_type);
+}
+
+template<typename T, typename C>
+__device__ __forceinline__ void welford_merge_block_vertical(C& count,
+                                                             T& mean,
+                                                             T& m2n,
+                                                             C* shmem_count,
+                                                             T* shmem_mean,
+                                                             T* shmem_m2n) {
+  // write to shared memory
+  auto address_base = threadIdx.x + threadIdx.y * blockDim.x;
+  shmem_mean[address_base] = mean;
+  shmem_m2n[address_base] = m2n;
+  shmem_count[address_base] = count;
+
+#pragma unroll
+  for (int offset = blockDim.y/2; offset > 0; offset >>= 1) {
+    __syncthreads();
+    if (threadIdx.y < offset && threadIdx.y + offset < blockDim.y) {
+      auto address = address_base + offset * blockDim.x;
+      // read shared memory back to register for reduction
+      auto num_new = shmem_count[address];
+      auto mean_new = shmem_mean[address];
+      auto m2n_new = shmem_m2n[address];
+
+      welford_merge_element(count, mean, m2n, num_new, mean_new, m2n_new);
+
+      // last write is not necessary
+      shmem_mean[address_base] = mean;
+      shmem_m2n[address_base] = m2n;
+      shmem_count[address_base] = count;
+    }
+  }
+}
+
+template<typename T>
+__device__ __forceinline__ void merge_block_vertical(T& sum_dy,
+                                                     T& sum_dy_xmu,
+                                                     T* shmem_sum_dy,
+                                                     T* shmem_sum_dy_xmu) {
+  // write to shared memory
+  auto address_base = threadIdx.x + threadIdx.y * blockDim.x;
+  shmem_sum_dy[address_base] = sum_dy;
+  shmem_sum_dy_xmu[address_base] = sum_dy_xmu;
+
+#pragma unroll
+  for (int offset = blockDim.y/2; offset > 0; offset >>= 1) {
+    __syncthreads();
+    if (threadIdx.y < offset && threadIdx.y + offset < blockDim.y) {
+      auto address = address_base + offset * blockDim.x;
+
+      sum_dy += shmem_sum_dy[address];
+      sum_dy_xmu += shmem_sum_dy_xmu[address];
+
+      // last write is not necessary
+      shmem_sum_dy[address_base] = sum_dy;
+      shmem_sum_dy_xmu[address_base] = sum_dy_xmu;
+    }
+  }
+}
+
+
+// welford kernel calculating mean/biased_variance/unbiased_variance
+template <typename scalar_t, typename accscalar_t, typename outscalar_t>
+__global__ void welford_kernel(
+      const scalar_t* __restrict__ input,
+      outscalar_t* __restrict__ out_mean,
+      outscalar_t* __restrict__ out_var_biased,
+      const int bs,
+      const int fs,
+      const int ss) {
+  int block_size = blockDim.x * blockDim.y;
+  int count = 0;
+  accscalar_t x_mean = accscalar_t(0);
+  accscalar_t m_2_n = accscalar_t(0);
+
+  int thread_id = threadIdx.y*blockDim.x + threadIdx.x;
+
+  for (int batch_id = threadIdx.y; batch_id < bs; batch_id += blockDim.y) {
+    int input_base = blockIdx.x*ss + batch_id*ss*fs;
+    // sequential welford
+    for (int offset = threadIdx.x; offset < ss ; offset += blockDim.x) {
+      count++;
+      auto x_n = static_cast<accscalar_t>(input[offset+input_base]);
+      auto d = x_n - x_mean;
+      x_mean += d / count;
+      m_2_n += d * (x_n - x_mean);
+    }
+  }
+
+  static __shared__ int s_mem[160];
+  accscalar_t* s_mem_ac = (accscalar_t*) &s_mem[32];
+
+  welford_reduce_mean_m2n<accscalar_t>(s_mem_ac, s_mem, x_mean, m_2_n, count, block_size, thread_id);
+
+  if (thread_id == 0) {
+    out_mean[blockIdx.x] = static_cast<outscalar_t>(x_mean);
+    out_var_biased[blockIdx.x] = static_cast<outscalar_t>(m_2_n/count);
+  }
+}
+
+// elementwise BN kernel
+template <typename scalar_t, typename accscalar_t, typename layerscalar_t>
+__global__ void batchnorm_forward_kernel(
+      const scalar_t* __restrict__ input,
+      const accscalar_t* __restrict__ mean,
+      const accscalar_t* __restrict__ inv_std,
+      const layerscalar_t* __restrict__ weight,
+      const layerscalar_t* __restrict__ shift,
+      scalar_t* __restrict__ out,
+      const int ss,
+      const int bs) {
+  auto m_c = mean[blockIdx.x];
+  auto inv_std_c = inv_std[blockIdx.x];
+  auto w_c = weight == NULL ? accscalar_t(1.0) : static_cast<accscalar_t>(weight[blockIdx.x]);
+  auto s_c = shift == NULL ? accscalar_t(0.0) : static_cast<accscalar_t>(shift[blockIdx.x]);
+
+  for (int batch_offset = blockIdx.y*blockDim.y + threadIdx.y; batch_offset < bs; batch_offset += gridDim.y*blockDim.y) {
+    int address_base = blockIdx.x*ss + batch_offset*gridDim.x*ss;
+    for (int offset = threadIdx.x + blockIdx.z*blockDim.x; offset < ss ; offset+= gridDim.z*blockDim.x) {
+      out[address_base+offset] = static_cast<scalar_t>(w_c * (static_cast<accscalar_t>(input[address_base+offset]) - m_c ) * inv_std_c + s_c);
+    }
+  }
+}
+
+// Backward BN kernel, calculates grad_bias, grad_weight as well as intermediate
+// results to calculating grad_input.
+// Breaking the grad_input to two step to support sync BN, which requires all
+// reduce of the intermediate results across processes.
+template <typename scalar_t, typename accscalar_t, typename layerscalar_t>
+__global__ void reduce_bn_kernel(
+      const scalar_t* __restrict__ input,
+      const scalar_t* __restrict__ grad_output,
+      const accscalar_t* __restrict__ mean,
+      const accscalar_t* __restrict__ inv_std,
+      accscalar_t* __restrict__ sum_dy_o,
+      accscalar_t* __restrict__ sum_dy_xmu_o,
+      layerscalar_t* __restrict__ grad_weight,
+      layerscalar_t* __restrict__ grad_bias,
+      const int bs,
+      const int fs,
+      const int ss) {
+  static __shared__ int s_mem[64];
+  //int total_item_num = bs * ss;
+
+  int thread_id = threadIdx.y*blockDim.x + threadIdx.x;
+
+  auto r_mean = mean[blockIdx.x];
+  auto factor = inv_std[blockIdx.x];
+
+  // Kahan sum
+  accscalar_t sum_dy = 0.0;
+  accscalar_t sum_dy_xmu = 0.0;
+  accscalar_t sum_dy_c = 0.0;
+  accscalar_t sum_dy_xmu_c = 0.0;
+  for (int batch_id = threadIdx.y; batch_id < bs; batch_id += blockDim.y) {
+    int input_base = blockIdx.x*ss + batch_id*ss*fs;
+    for (int offset = threadIdx.x; offset < ss ; offset += blockDim.x) {
+      auto e_grad = static_cast<accscalar_t>(grad_output[offset+input_base]);
+      auto e_input = static_cast<accscalar_t>(input[offset+input_base]);
+      // calculating sum_dy
+      auto sum_dy_y = e_grad - sum_dy_c;
+      auto sum_dy_t = sum_dy + sum_dy_y;
+      sum_dy_c = (sum_dy_t - sum_dy) - sum_dy_y;
+      sum_dy = sum_dy_t;
+
+      // calculating sum_dy_xmu
+      auto sum_dy_xmu_y = e_grad * (e_input - r_mean) - sum_dy_xmu_c;
+      auto sum_dy_xmu_t = sum_dy_xmu + sum_dy_xmu_y;
+      sum_dy_xmu_c = (sum_dy_xmu_t - sum_dy_xmu) - sum_dy_xmu_y;
+      sum_dy_xmu = sum_dy_xmu_t;
+    }
+  }
+
+  sum_dy = reduce_block((accscalar_t*)s_mem, sum_dy);
+  __syncthreads();
+  sum_dy_xmu = reduce_block((accscalar_t*)s_mem, sum_dy_xmu);
+
+  if (thread_id == 0) {
+    if (grad_bias != NULL) {
+      grad_bias[blockIdx.x] = static_cast<layerscalar_t>(sum_dy);
+    }
+    if (grad_weight != NULL) {
+      grad_weight[blockIdx.x] = static_cast<layerscalar_t>(sum_dy_xmu * factor);
+    }
+    //mean_dy[blockIdx.x] = sum_dy / total_item_num;
+    //mean_dy_xmu[blockIdx.x] = sum_dy_xmu / total_item_num;
+    sum_dy_o[blockIdx.x] = sum_dy;
+    sum_dy_xmu_o[blockIdx.x] = sum_dy_xmu;
+  }
+}
+
+// elementwise backward BN kernel
+template <typename scalar_t, typename accscalar_t, typename layerscalar_t>
+__global__ void batchnorm_backward_kernel(
+      const scalar_t* __restrict__ grad_output,
+      const scalar_t* __restrict__ input,
+      const accscalar_t* __restrict__ mean,
+      const accscalar_t* __restrict__ inv_std,
+      const layerscalar_t* __restrict__ weight,
+      const accscalar_t* __restrict__ sum_dy,
+      const accscalar_t* __restrict__ sum_dy_xmu,
+      const int* __restrict__ numel,
+      scalar_t* __restrict__ grad_input,
+      const int64_t world_size,
+      const int ss,
+      const int bs) {
+  int64_t div = 0;
+  for (int i = 0; i < world_size; i++) {
+    div += numel[i];
+  }
+  auto m_c = static_cast<accscalar_t>(mean[blockIdx.x]);
+  //auto m_dy_c = static_cast<accscalar_t>(mean_dy[blockIdx.x]);
+  auto m_dy_c = static_cast<accscalar_t>(sum_dy[blockIdx.x]) / div;
+  auto factor_1_c = inv_std[blockIdx.x];
+  auto factor_2_c = (weight == NULL ? accscalar_t(1.0) : static_cast<accscalar_t>(weight[blockIdx.x])) * factor_1_c;
+  //factor_1_c = factor_1_c * factor_1_c * mean_dy_xmu[blockIdx.x];
+  factor_1_c = factor_1_c * factor_1_c * sum_dy_xmu[blockIdx.x] / div;
+
+  for (int batch_offset = blockIdx.y*blockDim.y+threadIdx.y; batch_offset < bs; batch_offset += gridDim.y*blockDim.y) {
+    int address_base = blockIdx.x*ss + batch_offset*gridDim.x*ss;
+    for (int offset = threadIdx.x + blockIdx.z*blockDim.x; offset < ss ; offset+= gridDim.z*blockDim.x) {
+      grad_input[address_base+offset] = (static_cast<accscalar_t>(grad_output[address_base+offset]) - m_dy_c - (static_cast<accscalar_t>(input[address_base+offset]) - m_c) * factor_1_c) * factor_2_c;
+    }
+  }
+}
+
+// welford kernel for c last tensor calculating mean/biased_variance/unbiased_variance
+template
+   <typename scalar_t,
+    typename accscalar_t,
+    typename outscalar_t,
+    int PARALLEL_LOADS>
+__global__ void
+welford_kernel_c_last(
+      const scalar_t* __restrict__ input,
+      outscalar_t* __restrict__ out_mean,
+      outscalar_t* __restrict__ out_var_biased,
+      volatile accscalar_t* staging_data,
+      int* semaphores,
+      const int reduction_size,
+      const int stride) {
+  // hide latency with concurrency
+  accscalar_t x_mean[PARALLEL_LOADS];
+  accscalar_t m_2_n[PARALLEL_LOADS];
+  int count[PARALLEL_LOADS];
+
+#pragma unroll
+  for (int i = 0; i < PARALLEL_LOADS; i++) {
+    x_mean[i] = accscalar_t(0);
+    m_2_n[i] = accscalar_t(0);
+    count[i] = accscalar_t(0);
+  }
+  // tensor dimension (m,c)
+
+  // loop along m dimension
+  int inner_loop_stride = blockDim.y * gridDim.y;
+
+  // offset along m dimension
+  int m_offset = blockIdx.y * blockDim.y + threadIdx.y;
+  int c_offset = blockIdx.x * blockDim.x + threadIdx.x;
+
+  int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);
+  int address_base = m_offset * stride + c_offset;
+  int address_increment = inner_loop_stride * stride;
+
+  for (int i = 0; i < loop_count; i++) {
+    accscalar_t x_math[PARALLEL_LOADS];
+    accscalar_t x_count_inv[PARALLEL_LOADS];
+    accscalar_t is_valid[PARALLEL_LOADS];
+
+    // load multiple data in
+#pragma unroll
+    for (int j = 0; j < PARALLEL_LOADS; j++) {
+      if (c_offset < stride && m_offset < reduction_size) {
+        x_math[j] = input[address_base];
+        count[j]++;
+        x_count_inv[j] = accscalar_t(1) / count[j];
+        is_valid[j] = accscalar_t(1);
+      } else {
+        x_math[j] = accscalar_t(0);
+        x_count_inv[j] = accscalar_t(0);
+        is_valid[j] = accscalar_t(0);
+      }
+      m_offset += inner_loop_stride;
+      address_base += address_increment;
+    }
+
+    // calculate mean/m2n with welford
+#pragma unroll
+    for (int j = 0; j < PARALLEL_LOADS; j++) {
+      accscalar_t delta0 = x_math[j] - x_mean[j];
+      x_mean[j] += delta0 * x_count_inv[j];
+      accscalar_t delta1 = x_math[j] - x_mean[j];
+      m_2_n[j] += delta0 * delta1 * is_valid[j];
+    }
+  }
+
+  // thread reduction to accumulate mean/m_2_n/count between PARALLEL_LOADS
+#pragma unroll
+  for (int j = 1; j < PARALLEL_LOADS; j++) {
+    welford_merge_element(count[0], x_mean[0], m_2_n[0], count[j], x_mean[j], m_2_n[j]);
+  }
+
+  // release x_mean / m_2_n
+  auto mean_th = x_mean[0];
+  auto m2_th = m_2_n[0];
+  auto count_th = count[0];
+
+  // block-wise reduction with shared memory (since reduction cannot be done within a warp)
+  static __shared__ accscalar_t shmem_mean[MAX_BLOCK_SIZE];
+  static __shared__ accscalar_t shmem_m2n[MAX_BLOCK_SIZE];
+  static __shared__ int shmem_count[MAX_BLOCK_SIZE];
+
+  welford_merge_block_vertical(count_th, mean_th, m2_th, shmem_count, shmem_mean, shmem_m2n);
+
+  // grid reduction if needed (coop launch used at the first place)
+  if (gridDim.y > 1) {
+    volatile accscalar_t* staging_mean = staging_data;
+    volatile accscalar_t* staging_m2n = &staging_data[stride*gridDim.y];
+    volatile int* staging_count = reinterpret_cast<volatile int*>(&staging_m2n[stride*gridDim.y]);
+
+    address_base = c_offset + blockIdx.y * stride;
+    // write data to staging_data;
+    if (threadIdx.y == 0 && c_offset < stride) {
+      staging_mean[address_base] = mean_th;
+      staging_m2n[address_base] = m2_th;
+      staging_count[address_base] = count_th;
+    }
+
+    __threadfence();
+    __syncthreads(); // ensuring writes to staging_ is visible to all blocks
+
+    __shared__ bool is_last_block_done;
+    // mark block done
+    if (threadIdx.x == 0 && threadIdx.y == 0) {
+      int old = atomicAdd(&semaphores[blockIdx.x], 1);
+      is_last_block_done = (old == (gridDim.y-1));
+    }
+
+    __syncthreads();
+
+    // check that all data is now available in global memory
+    if (is_last_block_done) {
+      count_th = 0;
+      mean_th = accscalar_t(0.0);
+      m2_th = accscalar_t(0.0);
+
+      for (int y = threadIdx.y; y < gridDim.y; y += blockDim.y) {
+        address_base = c_offset + y * stride;
+        int num_new = c_offset < stride ? staging_count[address_base] : 0;
+        accscalar_t mean_new = c_offset < stride ? staging_mean[address_base] : accscalar_t(0.0);
+        accscalar_t m2n_new = c_offset < stride ? staging_m2n[address_base] : accscalar_t(0.0);
+
+        welford_merge_element(count_th, mean_th, m2_th, num_new, mean_new, m2n_new);
+      }
+
+      welford_merge_block_vertical(count_th, mean_th, m2_th, shmem_count, shmem_mean, shmem_m2n);
+      if (threadIdx.y == 0 && c_offset < stride) {
+        out_mean[c_offset] = static_cast<outscalar_t>(mean_th);
+        out_var_biased[c_offset] = static_cast<outscalar_t>(m2_th / count_th);
+      }
+    }
+  } else {
+    if (blockIdx.y == 0 && threadIdx.y == 0 && c_offset < stride) {
+      out_mean[c_offset] = static_cast<outscalar_t>(mean_th);
+      out_var_biased[c_offset] = static_cast<outscalar_t>(m2_th / count_th);
+    }
+  }
+}
+
+// parallel welford kernel to further reduce mean / biased_var
+// into mean / unbiased_var / inv_std across multiple processes.
+template <typename scalar_t>
+__global__ void welford_kernel_parallel(
+      const scalar_t* __restrict__ mean,
+      const scalar_t* __restrict__ var_biased,
+      const int* __restrict__ numel,
+      scalar_t* __restrict__ out_mean,
+      scalar_t* __restrict__ out_var,
+      scalar_t* __restrict__ inv_std,
+      const int world_size,
+      const int feature_size,
+      const float eps) {
+
+  for (int i = blockIdx.x * blockDim.x + threadIdx.x; i < feature_size; i += gridDim.x * blockDim.x) {
+    // load data;
+    int address = i;
+    scalar_t x_mean = 0;
+    scalar_t m_2_n = 0;
+    int count = 0;
+    for (int j = 0; j < world_size; j++) {
+      welford_merge_element(count, x_mean, m_2_n, numel[j], mean[address], var_biased[address]*numel[j]);
+      address += feature_size;
+    }
+    out_mean[i] = x_mean;
+    out_var[i] = m_2_n/ (count - 1);
+    inv_std[i] = scalar_t(1) / sqrt(m_2_n/count + eps);
+  }
+}
+
+// elementwise BN kernel
+template <
+    typename scalar_t,
+    typename accscalar_t,
+    typename layerscalar_t,
+    int PARALLEL_LOADS>
+__global__ void batchnorm_forward_c_last_kernel(
+      const scalar_t* __restrict__ input,
+      const scalar_t* __restrict__ z,
+      const accscalar_t* __restrict__ mean,
+      const accscalar_t* __restrict__ inv_std,
+      const layerscalar_t* __restrict__ weight,
+      const layerscalar_t* __restrict__ shift,
+      scalar_t* __restrict__ out,
+      const int reduction_size,
+      const int stride,
+      const bool fuse_relu) {
+  // tensor dimension (m,c)
+  // loop along m dimension
+  int inner_loop_stride = blockDim.y * gridDim.y;
+
+  // offset along m dimension
+  int m_offset = blockIdx.y * blockDim.y + threadIdx.y;
+  int c_offset = blockIdx.x * blockDim.x + threadIdx.x;
+
+  auto m_c = mean[c_offset];
+  auto inv_std_c = static_cast<accscalar_t>(inv_std[c_offset]);
+  auto w_c = weight == NULL ? accscalar_t(1.0) : static_cast<accscalar_t>(weight[c_offset]);
+  auto s_c = shift == NULL ? accscalar_t(0.0) : static_cast<accscalar_t>(shift[c_offset]);
+
+  int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);
+  int address_base = m_offset * stride + c_offset;
+  int address_increment = inner_loop_stride * stride;
+
+  for (int i = 0; i < loop_count; i++) {
+#pragma unroll
+    for (int j = 0; j < PARALLEL_LOADS; j++) {
+      if (c_offset < stride && m_offset < reduction_size) {
+        auto tmp = w_c * (static_cast<accscalar_t>(input[address_base]) - m_c ) * inv_std_c + s_c;
+        if (z != NULL) {
+          tmp += z[address_base];
+        }
+        out[address_base] = (fuse_relu && tmp <= accscalar_t(0.0) ? scalar_t(0.0) : static_cast<scalar_t>(tmp));
+      }
+      m_offset += inner_loop_stride;
+      address_base += address_increment;
+    }
+  }
+}
+
+// elementwise BN kernel
+template <
+    typename scalar_t,
+    typename accscalar_t,
+    typename layerscalar_t,
+    int PARALLEL_LOADS>
+__global__ void relu_backward_c_last_kernel(
+      const scalar_t* __restrict__ grad_output,
+      const scalar_t* __restrict__ input,
+      const scalar_t* __restrict__ z,
+      const accscalar_t* __restrict__ mean,
+      const accscalar_t* __restrict__ inv_std,
+      const layerscalar_t* __restrict__ weight,
+      const layerscalar_t* __restrict__ shift,
+      scalar_t* __restrict__ out,
+      const int reduction_size,
+      const int stride) {
+  // tensor dimension (m,c)
+  // loop along m dimension
+  int inner_loop_stride = blockDim.y * gridDim.y;
+
+  // offset along m dimension
+  int m_offset = blockIdx.y * blockDim.y + threadIdx.y;
+  int c_offset = blockIdx.x * blockDim.x + threadIdx.x;
+
+  auto m_c = mean[c_offset];
+  auto inv_std_c = static_cast<accscalar_t>(inv_std[c_offset]);
+  auto w_c = weight == NULL ? accscalar_t(1.0) : static_cast<accscalar_t>(weight[c_offset]);
+  auto s_c = shift == NULL ? accscalar_t(0.0) : static_cast<accscalar_t>(shift[c_offset]);
+
+  int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);
+  int address_base = m_offset * stride + c_offset;
+  int address_increment = inner_loop_stride * stride;
+
+  for (int i = 0; i < loop_count; i++) {
+#pragma unroll
+    for (int j = 0; j < PARALLEL_LOADS; j++) {
+      if (c_offset < stride && m_offset < reduction_size) {
+        auto tmp = w_c * (static_cast<accscalar_t>(input[address_base]) - m_c ) * inv_std_c + s_c;
+        if (z != NULL) {
+          tmp += z[address_base];
+        }
+        out[address_base] = (tmp <= accscalar_t(0.0) ? scalar_t(0.0) : grad_output[address_base]);
+      }
+      m_offset += inner_loop_stride;
+      address_base += address_increment;
+    }
+  }
+}
+
+// batchnorm backward kernel for c last tensor
+template
+   <typename scalar_t,
+    typename accscalar_t,
+    typename layerscalar_t,
+    int PARALLEL_LOADS>
+__global__ void reduce_bn_c_last_kernel(
+      const scalar_t* __restrict__ input,
+      const scalar_t* __restrict__ grad_output,
+      const accscalar_t* __restrict__ mean,
+      const accscalar_t* __restrict__ inv_std,
+      accscalar_t* __restrict__ sum_dy_o,
+      accscalar_t* __restrict__ sum_dy_xmu_o,
+      layerscalar_t* __restrict__ grad_weight,
+      layerscalar_t* __restrict__ grad_bias,
+      volatile accscalar_t* staging_data,
+      int* semaphores,
+      const int reduction_size,
+      const int stride) {
+
+  // hide latency with concurrency
+  accscalar_t sum_dy[PARALLEL_LOADS];
+  accscalar_t sum_dy_xmu[PARALLEL_LOADS];
+
+#pragma unroll
+  for (int i = 0; i < PARALLEL_LOADS; i++) {
+    sum_dy[i] = accscalar_t(0);
+    sum_dy_xmu[i] = accscalar_t(0);
+  }
+  // tensor dimension (m,c)
+
+  // loop along m dimension
+  int inner_loop_stride = blockDim.y * gridDim.y;
+
+  // offset along m dimension
+  int m_offset = blockIdx.y * blockDim.y + threadIdx.y;
+  int c_offset = blockIdx.x * blockDim.x + threadIdx.x;
+
+  int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);
+  int address_base = m_offset * stride + c_offset;
+  int address_increment = inner_loop_stride * stride;
+
+  auto r_mean = mean[c_offset];
+  auto factor = inv_std[c_offset];
+
+  for (int i = 0; i < loop_count; i++) {
+    accscalar_t x_input[PARALLEL_LOADS];
+    accscalar_t x_grad_output[PARALLEL_LOADS];
+
+    // load multiple data in
+#pragma unroll
+    for (int j = 0; j < PARALLEL_LOADS; j++) {
+      if (c_offset < stride && m_offset < reduction_size) {
+        x_input[j] = input[address_base];
+        x_grad_output[j] = grad_output[address_base];
+      } else {
+        x_input[j] = accscalar_t(0);
+        x_grad_output[j] = accscalar_t(0);
+      }
+      m_offset += inner_loop_stride;
+      address_base += address_increment;
+    }
+
+    // calculate sum_dy / sum_dy_xmu
+#pragma unroll
+    for (int j = 0; j < PARALLEL_LOADS; j++) {
+      sum_dy[j] += x_grad_output[j];
+      sum_dy_xmu[j] += x_grad_output[j] * (x_input[j] - r_mean);
+    }
+  }
+
+  // thread reduction to accumulate sum_dy / sum_dy_xmu between PARALLEL_LOADS
+#pragma unroll
+  for (int j = 1; j < PARALLEL_LOADS; j++) {
+    sum_dy[0] += sum_dy[j];
+    sum_dy_xmu[0] += sum_dy_xmu[j];
+  }
+
+  // release array of registers
+  auto sum_dy_th = sum_dy[0];
+  auto sum_dy_xmu_th = sum_dy_xmu[0];
+
+  // block-wise reduction with shared memory (since reduction cannot be done within a warp)
+  static __shared__ accscalar_t shmem_sum_dy[MAX_BLOCK_SIZE];
+  static __shared__ accscalar_t shmem_sum_dy_xmu[MAX_BLOCK_SIZE];
+
+  merge_block_vertical(sum_dy_th, sum_dy_xmu_th, shmem_sum_dy, shmem_sum_dy_xmu);
+
+  // grid reduction if needed (coop launch used at the first place)
+  if (gridDim.y > 1) {
+    volatile accscalar_t* staging_sum_dy = staging_data;
+    volatile accscalar_t* staging_sum_dy_xmu = &staging_data[stride*gridDim.y];
+
+    address_base = c_offset + blockIdx.y * stride;
+    // write data to staging_data;
+    if (threadIdx.y == 0 && c_offset < stride) {
+      staging_sum_dy[address_base] = sum_dy_th;
+      staging_sum_dy_xmu[address_base] = sum_dy_xmu_th;
+    }
+
+    __threadfence();
+    __syncthreads(); // ensuring writes to staging_ is visible to all blocks
+
+    __shared__ bool is_last_block_done;
+    // mark block done
+    if (threadIdx.x == 0 && threadIdx.y == 0) {
+      int old = atomicAdd(&semaphores[blockIdx.x], 1);
+      is_last_block_done = (old == (gridDim.y-1));
+    }
+
+    __syncthreads();
+
+    // check that all data is now available in global memory
+    if (is_last_block_done) {
+      sum_dy_th = accscalar_t(0.0);
+      sum_dy_xmu_th = accscalar_t(0.0);
+
+      for (int y = threadIdx.y; y < gridDim.y; y += blockDim.y) {
+        address_base = c_offset + y * stride;
+        sum_dy_th += (c_offset < stride ? staging_sum_dy[address_base] : accscalar_t(0.0));
+        sum_dy_xmu_th += (c_offset < stride ? staging_sum_dy_xmu[address_base] : accscalar_t(0.0));
+      }
+
+      merge_block_vertical(sum_dy_th, sum_dy_xmu_th, shmem_sum_dy, shmem_sum_dy_xmu);
+      if (threadIdx.y == 0 && c_offset < stride) {
+        if (grad_bias != NULL) {
+          grad_bias[c_offset] = static_cast<layerscalar_t>(sum_dy_th);
+        }
+        if (grad_weight != NULL) {
+          grad_weight[c_offset] = static_cast<layerscalar_t>(sum_dy_xmu_th * factor);
+        }
+        //mean_dy[c_offset] = sum_dy_th / reduction_size;
+        //mean_dy_xmu[c_offset] = sum_dy_xmu_th / reduction_size;
+        sum_dy_o[c_offset] = sum_dy_th;
+        sum_dy_xmu_o[c_offset] = sum_dy_xmu_th;
+      }
+    }
+  } else {
+    if (blockIdx.y == 0 && threadIdx.y == 0 && c_offset < stride) {
+      if (grad_bias != NULL) {
+        grad_bias[c_offset] = static_cast<layerscalar_t>(sum_dy_th);
+      }
+      if (grad_weight != NULL) {
+        grad_weight[c_offset] = static_cast<layerscalar_t>(sum_dy_xmu_th * factor);
+      }
+      //mean_dy[c_offset] = sum_dy_th / reduction_size;
+      //mean_dy_xmu[c_offset] = sum_dy_xmu_th / reduction_size;
+      sum_dy_o[c_offset] = sum_dy_th;
+      sum_dy_xmu_o[c_offset] = sum_dy_xmu_th;
+    }
+  }
+}
+
+// elementwise BN kernel
+template <
+    typename scalar_t,
+    typename accscalar_t,
+    typename layerscalar_t,
+    int PARALLEL_LOADS>
+__global__ void batchnorm_backward_c_last_kernel(
+      const scalar_t* __restrict__ grad_output,
+      const scalar_t* __restrict__ input,
+      const accscalar_t* __restrict__ mean,
+      const accscalar_t* __restrict__ inv_std,
+      const layerscalar_t* __restrict__ weight,
+      const accscalar_t* __restrict__ sum_dy,
+      const accscalar_t* __restrict__ sum_dy_xmu,
+      const int* __restrict__ numel,
+      scalar_t* __restrict__ grad_input,
+      const int64_t world_size,
+      const int reduction_size,
+      const int stride) {
+  int64_t div = 0;
+  for (int i = 0; i < world_size; i++) {
+    div += numel[i];
+  }
+  // tensor dimension (m,c)
+  // loop along m dimension
+  int inner_loop_stride = blockDim.y * gridDim.y;
+
+  // offset along m dimension
+  int m_offset = blockIdx.y * blockDim.y + threadIdx.y;
+  int c_offset = blockIdx.x * blockDim.x + threadIdx.x;
+
+  auto m_c = mean[c_offset];
+  auto m_dy_c = sum_dy[c_offset] / div;
+  auto factor_1_c = inv_std[c_offset];
+  auto factor_2_c = (weight == NULL? accscalar_t(1.0) : static_cast<accscalar_t>(weight[c_offset])) * factor_1_c;
+  factor_1_c = factor_1_c * factor_1_c * sum_dy_xmu[c_offset] / div;
+
+  int loop_count = 1 + (reduction_size - 1) / (inner_loop_stride * PARALLEL_LOADS);
+  int address_base = m_offset * stride + c_offset;
+  int address_increment = inner_loop_stride * stride;
+
+  for (int i = 0; i < loop_count; i++) {
+#pragma unroll
+    for (int j = 0; j < PARALLEL_LOADS; j++) {
+      if (c_offset < stride && m_offset < reduction_size) {
+        grad_input[address_base] = static_cast<scalar_t>(
+            (static_cast<accscalar_t>(grad_output[address_base]) - m_dy_c -
+            (static_cast<accscalar_t>(input[address_base]) - m_c) * factor_1_c)
+            * factor_2_c);
+      }
+      m_offset += inner_loop_stride;
+      address_base += address_increment;
+    }
+  }
+}
+
+std::vector<at::Tensor> welford_mean_var_CUDA(const at::Tensor input) {
+  const auto batch_size = input.size(0);
+  const auto feature_size = input.size(1);
+
+  auto space_size = get_tensor_spatial_size(input);
+  auto scalar_type = promote_scalartype(input);
+
+  at::Tensor out_var_biased = at::empty({feature_size}, input.options().dtype(scalar_type));
+  at::Tensor out_mean = at::empty({feature_size}, input.options().dtype(scalar_type));
+
+  int block_y = min(h_last_pow2(batch_size), int(MAX_BLOCK_SIZE / 32));
+  int block_x = max(1, min(MAX_BLOCK_SIZE / block_y, h_last_pow2(space_size)));
+  const dim3 block(block_x, block_y);
+  const dim3 grid(feature_size);
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "welford_mean_var_kernel",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      welford_kernel<scalar_t_0, accscalar_t, accscalar_t><<<grid, block, 0, stream>>>(
+          input.DATA_PTR<scalar_t_0>(),
+          out_mean.DATA_PTR<accscalar_t>(),
+          out_var_biased.DATA_PTR<accscalar_t>(),
+          batch_size,
+          feature_size,
+          space_size);
+    );
+  }
+
+  return {out_mean, out_var_biased};
+}
+
+at::Tensor batchnorm_forward_CUDA(
+    const at::Tensor input,
+    const at::Tensor mean,
+    const at::Tensor inv_std,
+    const at::optional<at::Tensor> weight,
+    const at::optional<at::Tensor> shift) {
+  const auto batch_size = input.size(0);
+  const auto feature_size = input.size(1);
+  at::Tensor out = at::empty_like(input);
+
+  auto space_size = get_tensor_spatial_size(input);
+
+  int block_x = max(32, min(MAX_BLOCK_SIZE, h_last_pow2(space_size)/4));
+  int block_y = max(1, min(MAX_BLOCK_SIZE/block_x, h_last_pow2(batch_size)/4));
+  const dim3 block(block_x, block_y);
+  int grid_z = max(1, min(65535, h_last_pow2(space_size)/4/block_x));
+  int batch_group_size = max(1, min(65535, h_last_pow2(batch_size)/block_y));
+  const dim3 grid(feature_size, batch_group_size, grid_z);
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  if (input.scalar_type() == at::ScalarType::Half
+      && weight.has_value() &&
+      weight.value().scalar_type() == at::ScalarType::Float) {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_forward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      batchnorm_forward_kernel<scalar_t_0, accscalar_t, accscalar_t><<<grid, block, 0, stream>>>(
+          input.DATA_PTR<scalar_t_0>(),
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? weight.value().DATA_PTR<accscalar_t>() : NULL,
+          shift.has_value() ? shift.value().DATA_PTR<accscalar_t>() : NULL,
+          out.DATA_PTR<scalar_t_0>(),
+          space_size,
+          batch_size);
+    );
+  } else {
+    if (weight.has_value()) {
+      TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),
+          "input.scalar_type() is not supported with weight.scalar_type()");
+    }
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_forward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      batchnorm_forward_kernel<scalar_t_0, accscalar_t, scalar_t_0><<<grid, block, 0, stream>>>(
+          input.DATA_PTR<scalar_t_0>(),
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? weight.value().DATA_PTR<scalar_t_0>() : NULL,
+          shift.has_value() ? shift.value().DATA_PTR<scalar_t_0>() : NULL,
+          out.DATA_PTR<scalar_t_0>(),
+          space_size,
+          batch_size);
+    );
+  }
+  return out;
+}
+
+std::vector<at::Tensor> reduce_bn_CUDA(
+    const at::Tensor grad_output,
+    const at::Tensor input,
+    const at::Tensor mean,
+    const at::Tensor inv_std,
+    const at::optional<at::Tensor> weight)
+{
+  const auto batch_size = input.size(0);
+  const auto feature_size = input.size(1);
+
+  auto scalar_type = promote_scalartype(input);
+
+  at::Tensor sum_dy = at::empty({feature_size}, mean.options());
+  at::Tensor sum_dy_xmu = at::empty({feature_size}, mean.options());
+
+  at::Tensor grad_weight;
+  at::Tensor grad_bias;
+  if (weight.has_value()) {
+    grad_weight = at::empty({feature_size}, weight.value().options());
+    grad_bias = at::empty({feature_size}, weight.value().options());
+  } else {
+    grad_weight = at::empty({0}, mean.options());
+    grad_bias = at::empty({0}, mean.options());
+  }
+
+  auto space_size = get_tensor_spatial_size(input);
+
+  int block_y = min(h_last_pow2(batch_size), int(MAX_BLOCK_SIZE/ 32));
+  int block_x = max(1, min(MAX_BLOCK_SIZE/ block_y, h_last_pow2(space_size)));
+  const dim3 block(block_x, block_y);
+  const dim3 grid(feature_size);
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  if (input.scalar_type() == at::ScalarType::Half
+      && weight.has_value() &&
+      weight.value().scalar_type() == at::ScalarType::Float) {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_backward_reduce",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      reduce_bn_kernel<scalar_t_0, accscalar_t, accscalar_t><<<grid, block, 0, stream>>>(
+          input.DATA_PTR<scalar_t_0>(),
+          grad_output.DATA_PTR<scalar_t_0>(),
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          sum_dy.DATA_PTR<accscalar_t>(),
+          sum_dy_xmu.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? grad_weight.DATA_PTR<accscalar_t>() : NULL,
+          weight.has_value() ? grad_bias.DATA_PTR<accscalar_t>() : NULL,
+          batch_size,
+          feature_size,
+          space_size);
+    );
+  } else {
+    if (weight.has_value()) {
+        TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),
+            "input.scalar_type() is not supported with weight.scalar_type()");
+    }
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_backward_reduce",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      reduce_bn_kernel<scalar_t_0, accscalar_t, scalar_t_0><<<grid, block, 0, stream>>>(
+          input.DATA_PTR<scalar_t_0>(),
+          grad_output.DATA_PTR<scalar_t_0>(),
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          sum_dy.DATA_PTR<accscalar_t>(),
+          sum_dy_xmu.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? grad_weight.DATA_PTR<scalar_t_0>() : NULL,
+          weight.has_value() ? grad_bias.DATA_PTR<scalar_t_0>() : NULL,
+          batch_size,
+          feature_size,
+          space_size);
+    );
+  }
+
+  return {sum_dy, sum_dy_xmu, grad_weight, grad_bias};
+}
+
+at::Tensor batchnorm_backward_CUDA(
+    const at::Tensor grad_output,
+    const at::Tensor input,
+    const at::Tensor mean,
+    const at::Tensor inv_std,
+    const at::optional<at::Tensor> weight,
+    const at::Tensor sum_dy,
+    const at::Tensor sum_dy_xmu,
+    const at::Tensor count) {
+  const auto batch_size = input.size(0);
+  const auto feature_size = input.size(1);
+
+  at::Tensor grad_input = at::empty_like(input);
+
+  auto space_size = get_tensor_spatial_size(input);
+
+  int block_x = max(32, min(MAX_BLOCK_SIZE, h_last_pow2(space_size)/4));
+  int block_y = max(1, min(MAX_BLOCK_SIZE/block_x, h_last_pow2(batch_size)/4));
+  const dim3 block(block_x, block_y);
+  int grid_z = max(1, min(65535, h_last_pow2(space_size)/4/block_x));
+  int batch_group_size = max(1, min(65535, h_last_pow2(batch_size)/block_y));
+  const dim3 grid(feature_size, batch_group_size, grid_z);
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  if (input.scalar_type() == at::ScalarType::Half
+      && weight.has_value() &&
+      weight.value().scalar_type() == at::ScalarType::Float) {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_backward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      batchnorm_backward_kernel<scalar_t_0, accscalar_t, accscalar_t><<<grid, block, 0, stream>>>(
+          grad_output.DATA_PTR<scalar_t_0>(),
+          input.DATA_PTR<scalar_t_0>(),
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? weight.value().DATA_PTR<accscalar_t>() : NULL,
+          sum_dy.DATA_PTR<accscalar_t>(),
+          sum_dy_xmu.DATA_PTR<accscalar_t>(),
+          count.DATA_PTR<int>(),
+          grad_input.DATA_PTR<scalar_t_0>(),
+          count.numel(),
+          space_size,
+          batch_size);
+    );
+  } else {
+    if (weight.has_value()) {
+      TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),
+          "input.scalar_type() is not supported with weight.scalar_type()");
+    }
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_backward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      batchnorm_backward_kernel<scalar_t_0, accscalar_t, scalar_t_0><<<grid, block, 0, stream>>>(
+          grad_output.DATA_PTR<scalar_t_0>(),
+          input.DATA_PTR<scalar_t_0>(),
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? weight.value().DATA_PTR<scalar_t_0>() : NULL,
+          sum_dy.DATA_PTR<accscalar_t>(),
+          sum_dy_xmu.DATA_PTR<accscalar_t>(),
+          count.DATA_PTR<int>(),
+          grad_input.DATA_PTR<scalar_t_0>(),
+          count.numel(),
+          space_size,
+          batch_size);
+    );
+  }
+
+  return grad_input;
+}
+
+std::vector<at::Tensor> welford_parallel_CUDA(const at::Tensor mean_feature_nodes,
+                                              const at::Tensor var_biased,
+                                              const at::Tensor numel,
+                                              const float eps) {
+  const auto world_size = mean_feature_nodes.size(0);
+  const auto feature_size = mean_feature_nodes.size(1);
+
+  at::Tensor out_var = at::empty({feature_size}, var_biased.options());
+  at::Tensor inv_std = at::empty_like(out_var);
+  at::Tensor out_mean = at::empty_like(out_var);
+
+  at::Tensor mean_feature_nodes_ = mean_feature_nodes.contiguous();
+  at::Tensor var_biased_ = var_biased.contiguous();
+  at::Tensor numel_ = numel.contiguous();
+
+  // TODO(jie): tile this for memory coalescing!
+  const int block = std::min(h_last_pow2(feature_size), MAX_BLOCK_SIZE);
+  const int grid = std::max<int>(1, feature_size / block);
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(mean_feature_nodes.scalar_type(), 0, "welford_parallel_kernel",
+      welford_kernel_parallel<scalar_t_0><<<grid, block, 0, stream>>>(
+          mean_feature_nodes_.DATA_PTR<scalar_t_0>(),
+          var_biased_.DATA_PTR<scalar_t_0>(),
+          numel_.DATA_PTR<int>(),
+          out_mean.DATA_PTR<scalar_t_0>(),
+          out_var.DATA_PTR<scalar_t_0>(),
+          inv_std.DATA_PTR<scalar_t_0>(),
+          world_size,
+          feature_size,
+          eps);
+    );
+  }
+
+  return {out_mean, out_var, inv_std};
+}
+
+std::vector<at::Tensor> welford_mean_var_c_last_CUDA(const at::Tensor input) {
+  const auto stride = input.size(input.ndimension()-1);
+  const auto reduction_size = input.numel() / stride;
+
+  auto scalar_type = promote_scalartype(input);
+  auto option = input.options().dtype(scalar_type);
+
+  at::Tensor out_var_biased = at::empty({stride}, option);
+  at::Tensor out_mean = at::empty({stride}, option);
+
+  dim3 block;
+  dim3 grid;
+  flexible_launch_configs(reduction_size, stride, block, grid, true);
+
+  at::Tensor staging_data;
+  at::Tensor semaphores;
+  if (grid.y > 1) {
+    staging_data = at::empty({4*stride*grid.y}, option);
+    semaphores = at::zeros({grid.x}, input.options().dtype(at::kInt));
+  }
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "welford_mean_var_c_last",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      accscalar_t* staging_data_ptr = grid.y > 1 ? staging_data.DATA_PTR<accscalar_t>() : nullptr;
+      int* semaphores_ptr = grid.y > 1 ? semaphores.DATA_PTR<int>() : nullptr;
+      welford_kernel_c_last<scalar_t_0, accscalar_t, accscalar_t, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          input.DATA_PTR<scalar_t_0>(),
+          out_mean.DATA_PTR<accscalar_t>(),
+          out_var_biased.DATA_PTR<accscalar_t>(),
+          staging_data_ptr,
+          semaphores_ptr,
+          reduction_size,
+          stride);
+    );
+  }
+
+  return {out_mean, out_var_biased};
+}
+
+at::Tensor batchnorm_forward_c_last_CUDA(
+    const at::Tensor input,
+    const at::optional<at::Tensor> z,
+    const at::Tensor mean,
+    const at::Tensor inv_std,
+    const at::optional<at::Tensor> weight,
+    const at::optional<at::Tensor> shift,
+    const bool fuse_relu) {
+  const auto stride = input.size(input.ndimension()-1);
+  const auto reduction_size = input.numel() / stride;
+
+  at::Tensor out = at::empty_like(input);
+
+  dim3 block;
+  dim3 grid;
+  flexible_launch_configs(reduction_size, stride, block, grid);
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  if (input.scalar_type() == at::ScalarType::Half
+      && weight.has_value() && weight.value().scalar_type() == at::ScalarType::Float) {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_forward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      batchnorm_forward_c_last_kernel<scalar_t_0, accscalar_t, accscalar_t, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          input.DATA_PTR<scalar_t_0>(),
+          z.has_value() ? z.value().DATA_PTR<scalar_t_0>() : NULL,
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? weight.value().DATA_PTR<accscalar_t>() : NULL,
+          shift.has_value() ? shift.value().DATA_PTR<accscalar_t>(): NULL,
+          out.DATA_PTR<scalar_t_0>(),
+          reduction_size,
+          stride,
+          fuse_relu);
+    );
+  } else {
+    if (weight.has_value()) {
+      TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),
+          "input.scalar_type() is not supported with weight.scalar_type()");
+    }
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_forward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      batchnorm_forward_c_last_kernel<scalar_t_0, accscalar_t, scalar_t_0, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          input.DATA_PTR<scalar_t_0>(),
+          z.has_value() ? z.value().DATA_PTR<scalar_t_0>() : NULL,
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? weight.value().DATA_PTR<scalar_t_0>() : NULL,
+          shift.has_value() ? shift.value().DATA_PTR<scalar_t_0>(): NULL,
+          out.DATA_PTR<scalar_t_0>(),
+          reduction_size,
+          stride,
+          fuse_relu);
+    );
+  }
+  return out;
+}
+
+std::vector<at::Tensor> reduce_bn_c_last_CUDA(
+    const at::Tensor grad_output,
+    const at::Tensor input,
+    const at::Tensor mean,
+    const at::Tensor inv_std,
+    const at::optional<at::Tensor> weight) {
+  const auto stride = input.size(input.ndimension()-1);
+  const auto reduction_size = input.numel() / stride;
+
+  at::Tensor sumn_dy = at::empty({stride}, mean.options());
+  at::Tensor sum_dy_xmu = at::empty({stride}, mean.options());
+
+  at::Tensor grad_weight;
+  at::Tensor grad_bias;
+  if (weight.has_value()) {
+    grad_weight = at::empty({stride}, weight.value().options());
+    grad_bias = at::empty({stride}, weight.value().options());
+  } else {
+    // because I cannot return an uninitialized at::Tensor
+    grad_weight = at::empty({0}, mean.options());
+    grad_bias = at::empty({0}, mean.options());
+  }
+
+  dim3 block;
+  dim3 grid;
+  flexible_launch_configs(reduction_size, stride, block, grid, true);
+
+  at::Tensor staging_data;
+  at::Tensor semaphores;
+  if (grid.y > 1) {
+    staging_data = at::empty({2*stride*grid.y}, mean.options());
+    semaphores = at::zeros({grid.x}, input.options().dtype(at::kInt));
+  }
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  if (input.scalar_type() == at::ScalarType::Half
+      && weight.has_value()
+      && weight.value().scalar_type() == at::ScalarType::Float) {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_backward_reduce",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      accscalar_t* staging_data_ptr = grid.y > 1 ? staging_data.DATA_PTR<accscalar_t>() : nullptr;
+      int* semaphores_ptr = grid.y > 1 ? semaphores.DATA_PTR<int>() : nullptr;
+      reduce_bn_c_last_kernel<scalar_t_0, accscalar_t, accscalar_t, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          input.DATA_PTR<scalar_t_0>(),
+          grad_output.DATA_PTR<scalar_t_0>(),
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          sumn_dy.DATA_PTR<accscalar_t>(),
+          sum_dy_xmu.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? grad_weight.DATA_PTR<accscalar_t>() : NULL,
+          weight.has_value() ?grad_bias.DATA_PTR<accscalar_t>() : NULL,
+          staging_data_ptr,
+          semaphores_ptr,
+          reduction_size,
+          stride);
+    );
+  } else {
+    if (weight.has_value()) {
+      TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),
+          "input.scalar_type() is not supported with weight.scalar_type()");
+    }
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_backward_reduce",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      accscalar_t* staging_data_ptr = grid.y > 1 ? staging_data.DATA_PTR<accscalar_t>() : nullptr;
+      int* semaphores_ptr = grid.y > 1 ? semaphores.DATA_PTR<int>() : nullptr;
+      reduce_bn_c_last_kernel<scalar_t_0, accscalar_t, scalar_t_0, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          input.DATA_PTR<scalar_t_0>(),
+          grad_output.DATA_PTR<scalar_t_0>(),
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          sumn_dy.DATA_PTR<accscalar_t>(),
+          sum_dy_xmu.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? grad_weight.DATA_PTR<scalar_t_0>() : NULL,
+          weight.has_value() ?grad_bias.DATA_PTR<scalar_t_0>() : NULL,
+          staging_data_ptr,
+          semaphores_ptr,
+          reduction_size,
+          stride);
+    );
+  }
+
+  return {sumn_dy, sum_dy_xmu, grad_weight, grad_bias};
+}
+
+at::Tensor batchnorm_backward_c_last_CUDA(
+    const at::Tensor grad_output,
+    const at::Tensor input,
+    const at::Tensor mean,
+    const at::Tensor inv_std,
+    const at::optional<at::Tensor> weight,
+    const at::Tensor sum_dy,
+    const at::Tensor sum_dy_xmu,
+    const at::Tensor count) {
+  const auto stride = input.size(input.ndimension()-1);
+  const auto reduction_size = input.numel() / stride;
+
+  at::Tensor grad_input = at::empty_like(input);
+
+  dim3 block;
+  dim3 grid;
+  flexible_launch_configs(reduction_size, stride, block, grid);
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  if (input.scalar_type() == at::ScalarType::Half
+      && weight.has_value() && weight.value().scalar_type() == at::ScalarType::Float) {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_forward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      batchnorm_backward_c_last_kernel<scalar_t_0, accscalar_t, accscalar_t, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          grad_output.DATA_PTR<scalar_t_0>(),
+          input.DATA_PTR<scalar_t_0>(),
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? weight.value().DATA_PTR<accscalar_t>() : NULL,
+          sum_dy.DATA_PTR<accscalar_t>(),
+          sum_dy_xmu.DATA_PTR<accscalar_t>(),
+          count.DATA_PTR<int>(),
+          grad_input.DATA_PTR<scalar_t_0>(),
+          count.numel(),
+          reduction_size,
+          stride);
+    );
+  } else {
+    if (weight.has_value()) {
+      TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),
+          "input.scalar_type() is not supported with weight.scalar_type()");
+    }
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_forward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      batchnorm_backward_c_last_kernel<scalar_t_0, accscalar_t, scalar_t_0, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          grad_output.DATA_PTR<scalar_t_0>(),
+          input.DATA_PTR<scalar_t_0>(),
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? weight.value().DATA_PTR<scalar_t_0>() : NULL,
+          sum_dy.DATA_PTR<accscalar_t>(),
+          sum_dy_xmu.DATA_PTR<accscalar_t>(),
+          count.DATA_PTR<int>(),
+          grad_input.DATA_PTR<scalar_t_0>(),
+          count.numel(),
+          reduction_size,
+          stride);
+    );
+  }
+ 
+  return grad_input;
+}
+
+at::Tensor relu_backward_c_last_CUDA(
+    const at::Tensor grad_output,
+    const at::Tensor input,
+    const at::optional<at::Tensor> z,
+    const at::Tensor mean,
+    const at::Tensor inv_std,
+    const at::optional<at::Tensor> weight,
+    const at::optional<at::Tensor> shift) {
+
+  const auto stride = input.size(input.ndimension()-1);
+  const auto reduction_size = input.numel() / stride;
+
+  at::Tensor out = at::empty_like(input);
+
+  dim3 block;
+  dim3 grid;
+  flexible_launch_configs(reduction_size, stride, block, grid);
+
+  auto stream = at::cuda::getCurrentCUDAStream();
+
+  if (input.scalar_type() == at::ScalarType::Half
+      && weight.has_value() && weight.value().scalar_type() == at::ScalarType::Float) {
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_forward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      relu_backward_c_last_kernel<scalar_t_0, accscalar_t, accscalar_t, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          grad_output.DATA_PTR<scalar_t_0>(),
+          input.DATA_PTR<scalar_t_0>(),
+          z.has_value() ? z.value().DATA_PTR<scalar_t_0>() : NULL,
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? weight.value().DATA_PTR<accscalar_t>() : NULL,
+          shift.has_value() ? shift.value().DATA_PTR<accscalar_t>(): NULL,
+          out.DATA_PTR<scalar_t_0>(),
+          reduction_size,
+          stride);
+    );
+  } else {
+    if (weight.has_value()) {
+      TORCH_CHECK(input.scalar_type() == weight.value().scalar_type(),
+          "input.scalar_type() is not supported with weight.scalar_type()");
+    }
+    using namespace at;
+    DISPATCH_FLOAT_AND_HALF(input.scalar_type(), 0, "batchnorm_forward",
+      using accscalar_t = at::acc_type<scalar_t_0, true>;
+      relu_backward_c_last_kernel<scalar_t_0, accscalar_t, scalar_t_0, ELEMENTS_PER_ITER>
+          <<<grid, block, 0, stream>>>(
+          grad_output.DATA_PTR<scalar_t_0>(),
+          input.DATA_PTR<scalar_t_0>(),
+          z.has_value() ? z.value().DATA_PTR<scalar_t_0>() : NULL,
+          mean.DATA_PTR<accscalar_t>(),
+          inv_std.DATA_PTR<accscalar_t>(),
+          weight.has_value() ? weight.value().DATA_PTR<scalar_t_0>() : NULL,
+          shift.has_value() ? shift.value().DATA_PTR<scalar_t_0>(): NULL,
+          out.DATA_PTR<scalar_t_0>(),
+          reduction_size,
+          stride);
+    );
+  }
+  return out;
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/Makefile b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/Makefile
new file mode 100644
index 0000000000..86cc24948d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/Makefile
@@ -0,0 +1,32 @@
+# Minimal makefile for Sphinx documentation
+#
+
+# You can set these variables from the command line.
+SPHINXOPTS    =
+SPHINXBUILD   = sphinx-build
+SPHINXPROJ    = NVIDIAAPEX
+SOURCEDIR     = source
+BUILDDIR      = build
+
+# Put it first so that "make" without argument is like "make help".
+help:
+	@$(SPHINXBUILD) -M help "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
+
+gh-pages:
+	git checkout gh-pages
+	rm -rf build
+	rm -rf source
+	git checkout master -- .
+	make html
+	rm -rf ../_modules ../_sources ../_static
+	mv -fv build/html/* ../
+	rm -rf build
+	git add -A
+	git commit -m "Generated gh-pages for `git log master -1 --pretty=short --abbrev-commit`" && git push origin gh-pages ; git checkout master
+
+.PHONY: help Makefile
+
+# Catch-all target: route all unknown targets to Sphinx using the new
+# "make mode" option.  $(O) is meant as a shortcut for $(SPHINXOPTS).
+%: Makefile
+	@$(SPHINXBUILD) -M $@ "$(SOURCEDIR)" "$(BUILDDIR)" $(SPHINXOPTS) $(O)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/_static/css/pytorch_theme.css b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/_static/css/pytorch_theme.css
new file mode 100644
index 0000000000..45e984c905
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/_static/css/pytorch_theme.css
@@ -0,0 +1,118 @@
+body {
+    font-family: "Lato","proxima-nova","Helvetica Neue",Arial,sans-serif;
+}
+
+/* Default header fonts are ugly */
+h1, h2, .rst-content .toctree-wrapper p.caption, h3, h4, h5, h6, legend, p.caption {
+    font-family: "Lato","proxima-nova","Helvetica Neue",Arial,sans-serif;
+}
+
+/* Use white for docs background */
+.wy-side-nav-search {
+    background-color: #fff;
+}
+
+.wy-nav-content-wrap, .wy-menu li.current > a  {
+    background-color: #fff;
+}
+
+@media screen and (min-width: 1400px) {
+    .wy-nav-content-wrap {
+        background-color: rgba(0, 0, 0, 0.0470588);
+    }
+
+    .wy-nav-content {
+        background-color: #fff;
+    }
+}
+
+/* Fixes for mobile */
+.wy-nav-top {
+    background-color: #fff;
+    background-image: url('../img/apex.jpg');
+    background-repeat: no-repeat;
+    background-position: center;
+    padding: 0;
+    margin: 0.4045em 0.809em;
+    color: #333;
+}
+
+.wy-nav-top > a {
+    display: none;
+}
+
+@media screen and (max-width: 768px) {
+    .wy-side-nav-search>a img.logo {
+        height: 60px;
+    }
+}
+
+/* This is needed to ensure that logo above search scales properly */
+.wy-side-nav-search a {
+    display: block;
+}
+
+/* This ensures that multiple constructors will remain in separate lines. */
+.rst-content dl:not(.docutils) dt {
+    display: table;
+}
+
+/* Use our red for literals (it's very similar to the original color) */
+.rst-content tt.literal, .rst-content tt.literal, .rst-content code.literal {
+    color: #F05732;
+}
+
+.rst-content tt.xref, a .rst-content tt, .rst-content tt.xref,
+.rst-content code.xref, a .rst-content tt, a .rst-content code {
+    color: #404040;
+}
+
+/* Change link colors (except for the menu) */
+
+a {
+    color: #F05732;
+}
+
+a:hover {
+    color: #F05732;
+}
+
+
+a:visited {
+    color: #D44D2C;
+}
+
+.wy-menu a {
+    color: #b3b3b3;
+}
+
+.wy-menu a:hover {
+    color: #b3b3b3;
+}
+
+/* Default footer text is quite big */
+footer {
+    font-size: 80%;
+}
+
+footer .rst-footer-buttons {
+    font-size: 125%; /* revert footer settings - 1/80% = 125% */
+}
+
+footer p {
+    font-size: 100%;
+}
+
+/* For hidden headers that appear in TOC tree */
+/* see http://stackoverflow.com/a/32363545/3343043 */
+.rst-content .hidden-section {
+    display: none;
+}
+
+nav .hidden-section {
+    display: inherit;
+}
+
+.wy-side-nav-search>div.version {
+    color: #000;
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/_templates/layout.html b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/_templates/layout.html
new file mode 100644
index 0000000000..63dfed9c94
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/_templates/layout.html
@@ -0,0 +1,51 @@
+{% extends "!layout.html" %}
+  {% block sidebartitle %} {{ super() }}
+
+  <style>
+    /* Sidebar header (and topbar for mobile) */
+    .wy-side-nav-search, .wy-nav-top {
+      background: #76b900;
+    }
+
+    .wy-side-nav-search a:link, .wy-nav-top a:link {
+      color: #fff;
+    }
+    .wy-side-nav-search a:visited, .wy-nav-top a:visited {
+      color: #fff;
+    }
+    .wy-side-nav-search a:hover, .wy-nav-top a:hover {
+      color: #fff;
+    }
+
+    .wy-menu-vertical a:link, .wy-menu-vertical a:visited {
+      color: #d9d9d9
+    }
+
+    .wy-menu-vertical a:active {
+      background-color: #76b900
+    }
+
+    .wy-side-nav-search>div.version {
+      color: rgba(0, 0, 0, 0.3)
+    }
+  </style>
+  {% endblock %}
+
+  {% block footer %} {{ super() }}
+
+  <style>
+  a:link, a:visited {
+    color: #76b900;
+  }
+
+  a:hover {
+    color: #8c0;
+  }
+
+  .rst-content dl:not(.docutils) dt {
+    background: rgba(118, 185, 0, 0.1);
+    color: rgba(59,93,0,1);
+    border-top: solid 3px rgba(59,93,0,1);
+  }
+  </style>
+  {% endblock %}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/conf.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/conf.py
new file mode 100644
index 0000000000..4477a28e13
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/docs/source/conf.py
@@ -0,0 +1,248 @@
+#!/usr/bin/env python3
+# -*- coding: utf-8 -*-
+#
+# PyTorch documentation build configuration file, created by
+# sphinx-quickstart on Fri Dec 23 13:31:47 2016.
+#
+# This file is execfile()d with the current directory set to its
+# containing dir.
+#
+# Note that not all possible configuration values are present in this
+# autogenerated file.
+#
+# All configuration values have a default; values that are commented out
+# serve to show the default.
+
+# If extensions (or modules to document with autodoc) are in another directory,
+# add these directories to sys.path here. If the directory is relative to the
+# documentation root, use os.path.abspath to make it absolute, like shown here.
+#
+import os
+import sys
+sys.path.insert(0, os.path.abspath('.'))
+# sys.path.insert(0, os.path.abspath('../../apex/parallel/'))
+import apex
+# import multiproc
+import sphinx_rtd_theme
+
+
+# -- General configuration ------------------------------------------------
+
+# If your documentation needs a minimal Sphinx version, state it here.
+#
+# needs_sphinx = '1.0'
+
+# Add any Sphinx extension module names here, as strings. They can be
+# extensions coming with Sphinx (named 'sphinx.ext.*') or your custom
+# ones.
+extensions = [
+    'sphinx.ext.autodoc',
+    'sphinx.ext.autosummary',
+    'sphinx.ext.doctest',
+    'sphinx.ext.intersphinx',
+    'sphinx.ext.todo',
+    'sphinx.ext.coverage',
+    'sphinx.ext.mathjax',
+    'sphinx.ext.napoleon',
+    'sphinx.ext.viewcode',
+    'sphinx.ext.extlinks',
+]
+
+napoleon_use_ivar = True
+
+# Add any paths that contain templates here, relative to this directory.
+templates_path = ['_templates']
+
+# The suffix(es) of source filenames.
+# You can specify multiple suffix as a list of string:
+#
+# source_suffix = ['.rst', '.md']
+source_suffix = '.rst'
+
+# The master toctree document.
+master_doc = 'index'
+
+# General information about the project.
+project = 'Apex'
+copyright = '2018'
+author = 'Christian Sarofeen, Natalia Gimelshein, Michael Carilli, Raul Puri'
+
+# The version info for the project you're documenting, acts as replacement for
+# |version| and |release|, also used in various other places throughout the
+# built documents.
+#
+# The short X.Y version.
+# TODO: change to [:2] at v1.0
+# version = 'master (' + torch.__version__ + ' )'
+version = '0.1'
+# The full version, including alpha/beta/rc tags.
+# TODO: verify this works as expected
+release = '0.1.0'
+
+# The language for content autogenerated by Sphinx. Refer to documentation
+# for a list of supported languages.
+#
+# This is also used if you do content translation via gettext catalogs.
+# Usually you set "language" from the command line for these cases.
+language = None
+
+# List of patterns, relative to source directory, that match files and
+# directories to ignore when looking for source files.
+# This pattern also affects html_static_path and html_extra_path
+exclude_patterns = []
+
+# The name of the Pygments (syntax highlighting) style to use.
+pygments_style = 'sphinx'
+
+# If true, `todo` and `todoList` produce output, else they produce nothing.
+todo_include_todos = True
+
+
+# -- Options for HTML output -------------------------------------------------
+
+# The theme to use for HTML and HTML Help pages.  See the documentation for
+# a list of builtin themes.
+#
+html_theme = 'sphinx_rtd_theme'
+html_theme_path = [sphinx_rtd_theme.get_html_theme_path()]
+
+# Theme options are theme-specific and customize the look and feel of a theme
+# further.  For a list of options available for each theme, see the
+# documentation.
+#
+html_theme_options = {
+    'collapse_navigation': False,
+    'display_version': True,
+    'logo_only': True,
+}
+
+# html_logo = '_static/img/nv-pytorch2.png'
+
+# Add any paths that contain custom static files (such as style sheets) here,
+# relative to this directory. They are copied after the builtin static files,
+# so a file named "default.css" will overwrite the builtin "default.css".
+html_static_path = ['_static']
+
+# html_style_path = 'css/pytorch_theme.css'
+html_context = {
+    'css_files': [
+        'https://fonts.googleapis.com/css?family=Lato',
+        '_static/css/pytorch_theme.css'
+    ],
+}
+
+
+# -- Options for HTMLHelp output ---------------------------------------------
+
+# Output file base name for HTML help builder.
+htmlhelp_basename = 'PyTorchdoc'
+
+
+# -- Options for LaTeX output ------------------------------------------------
+
+latex_elements = {
+    # The paper size ('letterpaper' or 'a4paper').
+    #
+    # 'papersize': 'letterpaper',
+
+    # The font size ('10pt', '11pt' or '12pt').
+    #
+    # 'pointsize': '10pt',
+
+    # Additional stuff for the LaTeX preamble.
+    #
+    # 'preamble': '',
+
+    # Latex figure (float) alignment
+    #
+    # 'figure_align': 'htbp',
+}
+
+# Grouping the document tree into LaTeX files. List of tuples
+# (source start file, target name, title,
+#  author, documentclass [howto, manual, or own class]).
+latex_documents = [
+    (master_doc, 'apex.tex', 'Apex Documentation',
+     'Torch Contributors', 'manual'),
+]
+
+
+# -- Options for manual page output ------------------------------------------
+
+# One entry per manual page. List of tuples
+# (source start file, name, description, authors, manual section).
+man_pages = [
+    (master_doc, 'Apex', 'Apex Documentation',
+     [author], 1)
+]
+
+
+# -- Options for Texinfo output ----------------------------------------------
+
+# Grouping the document tree into Texinfo files. List of tuples
+# (source start file, target name, title, author,
+#  dir menu entry, description, category)
+texinfo_documents = [
+    (master_doc, 'Apex', 'Apex Documentation',
+     author, 'Apex', 'One line description of project.',
+     'Miscellaneous'),
+]
+
+
+# Example configuration for intersphinx: refer to the Python standard library.
+intersphinx_mapping = {
+    'python': ('https://docs.python.org/', None),
+    'numpy': ('http://docs.scipy.org/doc/numpy/', None),
+}
+
+# -- A patch that prevents Sphinx from cross-referencing ivar tags -------
+# See http://stackoverflow.com/a/41184353/3343043
+
+from docutils import nodes
+from sphinx.util.docfields import TypedField
+from sphinx import addnodes
+
+
+def patched_make_field(self, types, domain, items, **kw):
+    # `kw` catches `env=None` needed for newer sphinx while maintaining
+    #  backwards compatibility when passed along further down!
+
+    # type: (List, unicode, Tuple) -> nodes.field
+    def handle_item(fieldarg, content):
+        par = nodes.paragraph()
+        par += addnodes.literal_strong('', fieldarg)  # Patch: this line added
+        # par.extend(self.make_xrefs(self.rolename, domain, fieldarg,
+        #                           addnodes.literal_strong))
+        if fieldarg in types:
+            par += nodes.Text(' (')
+            # NOTE: using .pop() here to prevent a single type node to be
+            # inserted twice into the doctree, which leads to
+            # inconsistencies later when references are resolved
+            fieldtype = types.pop(fieldarg)
+            if len(fieldtype) == 1 and isinstance(fieldtype[0], nodes.Text):
+                typename = u''.join(n.astext() for n in fieldtype)
+                typename = typename.replace('int', 'python:int')
+                typename = typename.replace('long', 'python:long')
+                typename = typename.replace('float', 'python:float')
+                typename = typename.replace('type', 'python:type')
+                par.extend(self.make_xrefs(self.typerolename, domain, typename,
+                                           addnodes.literal_emphasis, **kw))
+            else:
+                par += fieldtype
+            par += nodes.Text(')')
+        par += nodes.Text(' -- ')
+        par += content
+        return par
+
+    fieldname = nodes.field_name('', self.label)
+    if len(items) == 1 and self.can_collapse:
+        fieldarg, content = items[0]
+        bodynode = handle_item(fieldarg, content)
+    else:
+        bodynode = self.list_type()
+        for fieldarg, content in items:
+            bodynode += nodes.list_item('', handle_item(fieldarg, content))
+    fieldbody = nodes.field_body('', bodynode)
+    return nodes.field('', fieldname, fieldbody)
+
+TypedField.make_field = patched_make_field
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/README.md
new file mode 100644
index 0000000000..6cb9231813
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/README.md
@@ -0,0 +1,4 @@
+This directory contains examples illustrating Apex mixed precision and distributed tools.
+
+**Note for users of the pre-unification API**:
+`deprecated_api` contains examples illustrating the old (pre-unified) APIs.  These APIs will be removed soon, and users are strongly encouraged to switch.  The separate mixed precision tools called `Amp` and `FP16_Optimizer` in the old API are exposed via different flags/optimization levels in the new API.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/dcgan/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/dcgan/README.md
new file mode 100644
index 0000000000..9fc896cb5b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/dcgan/README.md
@@ -0,0 +1,41 @@
+# Mixed Precision DCGAN Training in PyTorch
+
+`main_amp.py` is based on [https://github.com/pytorch/examples/tree/master/dcgan](https://github.com/pytorch/examples/tree/master/dcgan).
+It implements Automatic Mixed Precision (Amp) training of the DCGAN example for different datasets. Command-line flags forwarded to `amp.initialize` are used to easily manipulate and switch between various pure and mixed precision "optimization levels" or `opt_level`s.  For a detailed explanation of `opt_level`s, see the [updated API guide](https://nvidia.github.io/apex/amp.html).
+
+We introduce these changes to the PyTorch DCGAN example as described in the [Multiple models/optimizers/losses](https://nvidia.github.io/apex/advanced.html#multiple-models-optimizers-losses) section of the documentation::
+```
+# Added after models and optimizers construction
+[netD, netG], [optimizerD, optimizerG] = amp.initialize(
+    [netD, netG], [optimizerD, optimizerG], opt_level=opt.opt_level, num_losses=3)
+...
+# loss.backward() changed to:
+with amp.scale_loss(errD_real, optimizerD, loss_id=0) as errD_real_scaled:
+    errD_real_scaled.backward()
+...
+with amp.scale_loss(errD_fake, optimizerD, loss_id=1) as errD_fake_scaled:
+    errD_fake_scaled.backward()
+...
+with amp.scale_loss(errG, optimizerG, loss_id=2) as errG_scaled:
+    errG_scaled.backward()
+```
+
+Note that we use different `loss_scalers` for each computed loss.
+Using a separate loss scaler per loss is [optional, not required](https://nvidia.github.io/apex/advanced.html#optionally-have-amp-use-a-different-loss-scaler-per-loss).
+
+To improve the numerical stability, we swapped `nn.Sigmoid() + nn.BCELoss()` to `nn.BCEWithLogitsLoss()`.
+
+With the new Amp API **you never need to explicitly convert your model, or the input data, to half().**
+
+"Pure FP32" training:
+```
+$ python main_amp.py --opt_level O0
+```
+Recommended mixed precision training:
+```
+$ python main_amp.py --opt_level O1
+```
+
+Have a look at the original [DCGAN example](https://github.com/pytorch/examples/tree/master/dcgan) for more information about the used arguments.
+
+To enable mixed precision training, we introduce the `--opt_level` argument.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/dcgan/main_amp.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/dcgan/main_amp.py
new file mode 100644
index 0000000000..be1a2894fd
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/dcgan/main_amp.py
@@ -0,0 +1,274 @@
+from __future__ import print_function
+import argparse
+import os
+import random
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.backends.cudnn as cudnn
+import torch.optim as optim
+import torch.utils.data
+import torchvision.datasets as dset
+import torchvision.transforms as transforms
+import torchvision.utils as vutils
+
+try:    
+    from apex import amp
+except ImportError:
+    raise ImportError("Please install apex from https://www.github.com/nvidia/apex to run this example.")
+
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--dataset', default='cifar10', help='cifar10 | lsun | mnist |imagenet | folder | lfw | fake')
+parser.add_argument('--dataroot', default='./', help='path to dataset')
+parser.add_argument('--workers', type=int, help='number of data loading workers', default=2)
+parser.add_argument('--batchSize', type=int, default=64, help='input batch size')
+parser.add_argument('--imageSize', type=int, default=64, help='the height / width of the input image to network')
+parser.add_argument('--nz', type=int, default=100, help='size of the latent z vector')
+parser.add_argument('--ngf', type=int, default=64)
+parser.add_argument('--ndf', type=int, default=64)
+parser.add_argument('--niter', type=int, default=25, help='number of epochs to train for')
+parser.add_argument('--lr', type=float, default=0.0002, help='learning rate, default=0.0002')
+parser.add_argument('--beta1', type=float, default=0.5, help='beta1 for adam. default=0.5')
+parser.add_argument('--ngpu', type=int, default=1, help='number of GPUs to use')
+parser.add_argument('--netG', default='', help="path to netG (to continue training)")
+parser.add_argument('--netD', default='', help="path to netD (to continue training)")
+parser.add_argument('--outf', default='.', help='folder to output images and model checkpoints')
+parser.add_argument('--manualSeed', type=int, help='manual seed')
+parser.add_argument('--classes', default='bedroom', help='comma separated list of classes for the lsun data set')
+parser.add_argument('--opt_level', default='O1', help='amp opt_level, default="O1"')
+
+opt = parser.parse_args()
+print(opt)
+
+
+try:
+    os.makedirs(opt.outf)
+except OSError:
+    pass
+
+if opt.manualSeed is None:
+    opt.manualSeed = 2809
+print("Random Seed: ", opt.manualSeed)
+random.seed(opt.manualSeed)
+torch.manual_seed(opt.manualSeed)
+
+cudnn.benchmark = True
+
+
+if opt.dataset in ['imagenet', 'folder', 'lfw']:
+    # folder dataset
+    dataset = dset.ImageFolder(root=opt.dataroot,
+                               transform=transforms.Compose([
+                                   transforms.Resize(opt.imageSize),
+                                   transforms.CenterCrop(opt.imageSize),
+                                   transforms.ToTensor(),
+                                   transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+                               ]))
+    nc=3
+elif opt.dataset == 'lsun':
+    classes = [ c + '_train' for c in opt.classes.split(',')]
+    dataset = dset.LSUN(root=opt.dataroot, classes=classes,
+                        transform=transforms.Compose([
+                            transforms.Resize(opt.imageSize),
+                            transforms.CenterCrop(opt.imageSize),
+                            transforms.ToTensor(),
+                            transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+                        ]))
+    nc=3
+elif opt.dataset == 'cifar10':
+    dataset = dset.CIFAR10(root=opt.dataroot, download=True,
+                           transform=transforms.Compose([
+                               transforms.Resize(opt.imageSize),
+                               transforms.ToTensor(),
+                               transforms.Normalize((0.5, 0.5, 0.5), (0.5, 0.5, 0.5)),
+                           ]))
+    nc=3
+
+elif opt.dataset == 'mnist':
+        dataset = dset.MNIST(root=opt.dataroot, download=True,
+                           transform=transforms.Compose([
+                               transforms.Resize(opt.imageSize),
+                               transforms.ToTensor(),
+                               transforms.Normalize((0.5,), (0.5,)),
+                           ]))
+        nc=1
+
+elif opt.dataset == 'fake':
+    dataset = dset.FakeData(image_size=(3, opt.imageSize, opt.imageSize),
+                            transform=transforms.ToTensor())
+    nc=3
+
+assert dataset
+dataloader = torch.utils.data.DataLoader(dataset, batch_size=opt.batchSize,
+                                         shuffle=True, num_workers=int(opt.workers))
+
+device = torch.device("cuda:0")
+ngpu = int(opt.ngpu)
+nz = int(opt.nz)
+ngf = int(opt.ngf)
+ndf = int(opt.ndf)
+
+
+# custom weights initialization called on netG and netD
+def weights_init(m):
+    classname = m.__class__.__name__
+    if classname.find('Conv') != -1:
+        m.weight.data.normal_(0.0, 0.02)
+    elif classname.find('BatchNorm') != -1:
+        m.weight.data.normal_(1.0, 0.02)
+        m.bias.data.fill_(0)
+
+
+class Generator(nn.Module):
+    def __init__(self, ngpu):
+        super(Generator, self).__init__()
+        self.ngpu = ngpu
+        self.main = nn.Sequential(
+            # input is Z, going into a convolution
+            nn.ConvTranspose2d(     nz, ngf * 8, 4, 1, 0, bias=False),
+            nn.BatchNorm2d(ngf * 8),
+            nn.ReLU(True),
+            # state size. (ngf*8) x 4 x 4
+            nn.ConvTranspose2d(ngf * 8, ngf * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf * 4),
+            nn.ReLU(True),
+            # state size. (ngf*4) x 8 x 8
+            nn.ConvTranspose2d(ngf * 4, ngf * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf * 2),
+            nn.ReLU(True),
+            # state size. (ngf*2) x 16 x 16
+            nn.ConvTranspose2d(ngf * 2,     ngf, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ngf),
+            nn.ReLU(True),
+            # state size. (ngf) x 32 x 32
+            nn.ConvTranspose2d(    ngf,      nc, 4, 2, 1, bias=False),
+            nn.Tanh()
+            # state size. (nc) x 64 x 64
+        )
+
+    def forward(self, input):
+        if input.is_cuda and self.ngpu > 1:
+            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
+        else:
+            output = self.main(input)
+        return output
+
+
+netG = Generator(ngpu).to(device)
+netG.apply(weights_init)
+if opt.netG != '':
+    netG.load_state_dict(torch.load(opt.netG))
+print(netG)
+
+
+class Discriminator(nn.Module):
+    def __init__(self, ngpu):
+        super(Discriminator, self).__init__()
+        self.ngpu = ngpu
+        self.main = nn.Sequential(
+            # input is (nc) x 64 x 64
+            nn.Conv2d(nc, ndf, 4, 2, 1, bias=False),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (ndf) x 32 x 32
+            nn.Conv2d(ndf, ndf * 2, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 2),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (ndf*2) x 16 x 16
+            nn.Conv2d(ndf * 2, ndf * 4, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 4),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (ndf*4) x 8 x 8
+            nn.Conv2d(ndf * 4, ndf * 8, 4, 2, 1, bias=False),
+            nn.BatchNorm2d(ndf * 8),
+            nn.LeakyReLU(0.2, inplace=True),
+            # state size. (ndf*8) x 4 x 4
+            nn.Conv2d(ndf * 8, 1, 4, 1, 0, bias=False),
+        )
+
+    def forward(self, input):
+        if input.is_cuda and self.ngpu > 1:
+            output = nn.parallel.data_parallel(self.main, input, range(self.ngpu))
+        else:
+            output = self.main(input)
+
+        return output.view(-1, 1).squeeze(1)
+
+
+netD = Discriminator(ngpu).to(device)
+netD.apply(weights_init)
+if opt.netD != '':
+    netD.load_state_dict(torch.load(opt.netD))
+print(netD)
+
+criterion = nn.BCEWithLogitsLoss()
+
+fixed_noise = torch.randn(opt.batchSize, nz, 1, 1, device=device)
+real_label = 1
+fake_label = 0
+
+# setup optimizer
+optimizerD = optim.Adam(netD.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
+optimizerG = optim.Adam(netG.parameters(), lr=opt.lr, betas=(opt.beta1, 0.999))
+
+[netD, netG], [optimizerD, optimizerG] = amp.initialize(
+    [netD, netG], [optimizerD, optimizerG], opt_level=opt.opt_level, num_losses=3)
+
+for epoch in range(opt.niter):
+    for i, data in enumerate(dataloader, 0):
+        ############################
+        # (1) Update D network: maximize log(D(x)) + log(1 - D(G(z)))
+        ###########################
+        # train with real
+        netD.zero_grad()
+        real_cpu = data[0].to(device)
+        batch_size = real_cpu.size(0)
+        label = torch.full((batch_size,), real_label, device=device)
+
+        output = netD(real_cpu)
+        errD_real = criterion(output, label)
+        with amp.scale_loss(errD_real, optimizerD, loss_id=0) as errD_real_scaled:
+            errD_real_scaled.backward()
+        D_x = output.mean().item()
+
+        # train with fake
+        noise = torch.randn(batch_size, nz, 1, 1, device=device)
+        fake = netG(noise)
+        label.fill_(fake_label)
+        output = netD(fake.detach())
+        errD_fake = criterion(output, label)
+        with amp.scale_loss(errD_fake, optimizerD, loss_id=1) as errD_fake_scaled:
+            errD_fake_scaled.backward()
+        D_G_z1 = output.mean().item()
+        errD = errD_real + errD_fake
+        optimizerD.step()
+
+        ############################
+        # (2) Update G network: maximize log(D(G(z)))
+        ###########################
+        netG.zero_grad()
+        label.fill_(real_label)  # fake labels are real for generator cost
+        output = netD(fake)
+        errG = criterion(output, label)
+        with amp.scale_loss(errG, optimizerG, loss_id=2) as errG_scaled:
+            errG_scaled.backward()
+        D_G_z2 = output.mean().item()
+        optimizerG.step()
+
+        print('[%d/%d][%d/%d] Loss_D: %.4f Loss_G: %.4f D(x): %.4f D(G(z)): %.4f / %.4f'
+              % (epoch, opt.niter, i, len(dataloader),
+                 errD.item(), errG.item(), D_x, D_G_z1, D_G_z2))
+        if i % 100 == 0:
+            vutils.save_image(real_cpu,
+                    '%s/real_samples.png' % opt.outf,
+                    normalize=True)
+            fake = netG(fixed_noise)
+            vutils.save_image(fake.detach(),
+                    '%s/amp_fake_samples_epoch_%03d.png' % (opt.outf, epoch),
+                    normalize=True)
+
+    # do checkpointing
+    torch.save(netG.state_dict(), '%s/netG_epoch_%d.pth' % (opt.outf, epoch))
+    torch.save(netD.state_dict(), '%s/netD_epoch_%d.pth' % (opt.outf, epoch))
+
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/docker/Dockerfile b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/docker/Dockerfile
new file mode 100644
index 0000000000..47fd422cdd
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/docker/Dockerfile
@@ -0,0 +1,16 @@
+# Base image must at least have pytorch and CUDA installed.
+ARG BASE_IMAGE=nvcr.io/nvidia/pytorch:19.07-py3
+FROM $BASE_IMAGE
+ARG BASE_IMAGE
+RUN echo "Installing Apex on top of ${BASE_IMAGE}"
+# make sure we don't overwrite some existing directory called "apex"
+WORKDIR /tmp/unique_for_apex
+# uninstall Apex if present, twice to make absolutely sure :)
+RUN pip uninstall -y apex || :
+RUN pip uninstall -y apex || :
+# SHA is something the user can touch to force recreation of this Docker layer,
+# and therefore force cloning of the latest version of Apex
+RUN SHA=ToUcHMe git clone https://github.com/NVIDIA/apex.git
+WORKDIR /tmp/unique_for_apex/apex
+RUN pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .
+WORKDIR /workspace
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/docker/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/docker/README.md
new file mode 100644
index 0000000000..3969af61a6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/docker/README.md
@@ -0,0 +1,40 @@
+## Option 1:  Create a new container with Apex
+
+**Dockerfile** installs the latest Apex on top of an existing image.  Run
+```
+docker build -t new_image_with_apex .
+```
+By default, **Dockerfile** uses NVIDIA's Pytorch container as the base image,
+which requires an NVIDIA GPU Cloud (NGC) account.  If you don't have an NGC account, you can sign up for free by following the instructions [here](https://docs.nvidia.com/ngc/ngc-getting-started-guide/index.html#generating-api-key).
+
+Alternatively, you can supply your own base image via the `BASE_IMAGE` build-arg.
+`BASE_IMAGE` must have Pytorch and Cuda installed.  For example, any
+`-devel` image for Pytorch 1.0 and later from the
+[official Pytorch Dockerhub](https://hub.docker.com/r/pytorch/pytorch) may be used:
+```
+docker build --build-arg BASE_IMAGE=1.3-cuda10.1-cudnn7-devel -t new_image_with_apex .
+```
+
+If you want to rebuild your image, and force the latest Apex to be cloned and installed, make any small change to the `SHA` variable in **Dockerfile**.
+
+**Warning:**
+Currently, the non-`-devel` images on Pytorch Dockerhub do not contain the Cuda compiler `nvcc`.  Therefore,
+images whose name does not contain `-devel` are not eligible candidates for `BASE_IMAGE`.
+
+### Running your Apex container
+
+Like any Cuda-enabled Pytorch container, a container with Apex should be run via [nvidia-docker](https://github.com/NVIDIA/nvidia-docker), for example:
+```
+docker run --runtime=nvidia -it --rm --ipc=host new_image_with_apex
+```
+
+## Option 2:  Install Apex in a running container
+
+Instead of building a new container, it is also a viable option to `git clone https://github.com/NVIDIA/apex.git` on bare metal, mount the Apex repo into your container at launch by running, for example,
+```
+docker run --runtime=nvidia -it --rm --ipc=host -v /bare/metal/apex:/apex/in/container <base image>
+```
+then go to /apex/in/container within the running container and
+```
+pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .
+```
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/imagenet/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/imagenet/README.md
new file mode 100644
index 0000000000..257d4a78d2
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/imagenet/README.md
@@ -0,0 +1,183 @@
+# Mixed Precision ImageNet Training in PyTorch
+
+`main_amp.py` is based on [https://github.com/pytorch/examples/tree/master/imagenet](https://github.com/pytorch/examples/tree/master/imagenet).
+It implements Automatic Mixed Precision (Amp) training of popular model architectures, such as ResNet, AlexNet, and VGG, on the ImageNet dataset.  Command-line flags forwarded to `amp.initialize` are used to easily manipulate and switch between various pure and mixed precision "optimization levels" or `opt_level`s.  For a detailed explanation of `opt_level`s, see the [updated API guide](https://nvidia.github.io/apex/amp.html).
+
+Three lines enable Amp:
+```
+# Added after model and optimizer construction
+model, optimizer = amp.initialize(model, optimizer, flags...)
+...
+# loss.backward() changed to:
+with amp.scale_loss(loss, optimizer) as scaled_loss:
+    scaled_loss.backward()
+```
+
+With the new Amp API **you never need to explicitly convert your model, or the input data, to half().**
+
+## Requirements
+
+- Download the ImageNet dataset and move validation images to labeled subfolders
+    - The following script may be helpful: https://raw.githubusercontent.com/soumith/imagenetloader.torch/master/valprep.sh
+
+## Training
+
+To train a model, create softlinks to the Imagenet dataset, then run `main.py` with the desired model architecture, as shown in `Example commands` below.
+
+The default learning rate schedule is set for ResNet50.  `main_amp.py` script rescales the learning rate according to the global batch size (number of distributed processes \* per-process minibatch size).
+
+## Example commands
+
+**Note:**  batch size `--b 224` assumes your GPUs have >=16GB of onboard memory.  You may be able to increase this to 256, but that's cutting it close, so it may out-of-memory for different Pytorch versions.
+
+**Note:**  All of the following use 4 dataloader subprocesses (`--workers 4`) to reduce potential
+CPU data loading bottlenecks.
+
+**Note:**  `--opt-level` `O1` and `O2` both use dynamic loss scaling by default unless manually overridden.
+`--opt-level` `O0` and `O3` (the "pure" training modes) do not use loss scaling by default.
+`O0` and `O3` can be told to use loss scaling via manual overrides, but using loss scaling with `O0`
+(pure FP32 training) does not really make sense, and will trigger a warning.
+
+Softlink training and validation datasets into the current directory:
+```
+$ ln -sf /data/imagenet/train-jpeg/ train
+$ ln -sf /data/imagenet/val-jpeg/ val
+```
+
+### Summary
+
+Amp allows easy experimentation with various pure and mixed precision options.
+```
+$ python main_amp.py -a resnet50 --b 128 --workers 4 --opt-level O0 ./
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O3 ./
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O3 --keep-batchnorm-fp32 True ./
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 ./
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 --loss-scale 128.0 ./
+$ python -m torch.distributed.launch --nproc_per_node=2 main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 ./
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 ./
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 --loss-scale 128.0 ./
+$ python -m torch.distributed.launch --nproc_per_node=2 main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 ./
+```
+Options are explained below.  Again, the [updated API guide](https://nvidia.github.io/apex/amp.html) provides more detail.
+
+#### `--opt-level O0` (FP32 training) and `O3` (FP16 training)
+
+"Pure FP32" training:
+```
+$ python main_amp.py -a resnet50 --b 128 --workers 4 --opt-level O0 ./
+```
+"Pure FP16" training:
+```
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O3 ./
+```
+FP16 training with FP32 batchnorm:
+```
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O3 --keep-batchnorm-fp32 True ./
+```
+Keeping the batchnorms in FP32 improves stability and allows Pytorch
+to use cudnn batchnorms, which significantly increases speed in Resnet50.
+
+The `O3` options might not converge, because they are not true mixed precision.
+However, they can be useful to establish "speed of light" performance for
+your model, which provides a baseline for comparison with `O1` and `O2`.
+For Resnet50 in particular, `--opt-level O3 --keep-batchnorm-fp32 True` establishes
+the "speed of light."  (Without `--keep-batchnorm-fp32`, it's slower, because it does
+not use cudnn batchnorm.)
+
+#### `--opt-level O1` (Official Mixed Precision recipe, recommended for typical use)
+
+`O1` patches Torch functions to cast inputs according to a whitelist-blacklist model.
+FP16-friendly (Tensor Core) ops like gemms and convolutions run in FP16, while ops
+that benefit from FP32, like batchnorm and softmax, run in FP32.
+Also, dynamic loss scaling is used by default.
+```
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 ./
+```
+`O1` overridden to use static loss scaling:
+```
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 --loss-scale 128.0
+```
+Distributed training with 2 processes (1 GPU per process, see **Distributed training** below
+for more detail)
+```
+$ python -m torch.distributed.launch --nproc_per_node=2 main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O1 ./
+```
+For best performance, set `--nproc_per_node` equal to the total number of GPUs on the node
+to use all available resources.
+
+#### `--opt-level O2` ("Almost FP16" mixed precision.  More dangerous than O1.)
+
+`O2` exists mainly to support some internal use cases.  Please prefer `O1`.
+
+`O2` casts the model to FP16, keeps batchnorms in FP32,
+maintains master weights in FP32, and implements
+dynamic loss scaling by default. (Unlike --opt-level O1, --opt-level O2
+does not patch Torch functions.)
+```
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 ./
+```
+"Fast mixed precision" overridden to use static loss scaling:
+```
+$ python main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 --loss-scale 128.0 ./
+```
+Distributed training with 2 processes (1 GPU per process)
+```
+$ python -m torch.distributed.launch --nproc_per_node=2 main_amp.py -a resnet50 --b 224 --workers 4 --opt-level O2 ./
+```
+
+## Distributed training
+
+`main_amp.py` optionally uses `apex.parallel.DistributedDataParallel` (DDP) for multiprocess training with one GPU per process.
+```
+model = apex.parallel.DistributedDataParallel(model)
+```
+is a drop-in replacement for
+```
+model = torch.nn.parallel.DistributedDataParallel(model,
+                                                  device_ids=[arg.local_rank],
+                                                  output_device=arg.local_rank)
+```
+(because Torch DDP permits multiple GPUs per process, with Torch DDP you are required to
+manually specify the device to run on and the output device.
+With Apex DDP, it uses only the current device by default).
+
+The choice of DDP wrapper (Torch or Apex) is orthogonal to the use of Amp and other Apex tools.  It is safe to use `apex.amp` with either `torch.nn.parallel.DistributedDataParallel` or `apex.parallel.DistributedDataParallel`.  In the future, I may add some features that permit optional tighter integration between `Amp` and `apex.parallel.DistributedDataParallel` for marginal performance benefits, but currently, there's no compelling reason to use Apex DDP versus Torch DDP for most models.
+
+To use DDP with `apex.amp`, the only gotcha is that
+```
+model, optimizer = amp.initialize(model, optimizer, flags...)
+```
+must precede
+```
+model = DDP(model)
+```
+If DDP wrapping occurs before `amp.initialize`, `amp.initialize` will raise an error.
+
+With both Apex DDP and Torch DDP, you must also call `torch.cuda.set_device(args.local_rank)` within
+each process prior to initializing your model or any other tensors.
+More information can be found in the docs for the
+Pytorch multiprocess launcher module [torch.distributed.launch](https://pytorch.org/docs/stable/distributed.html#launch-utility).
+
+`main_amp.py` is written to interact with 
+[torch.distributed.launch](https://pytorch.org/docs/master/distributed.html#launch-utility),
+which spawns multiprocess jobs using the following syntax:
+```
+python -m torch.distributed.launch --nproc_per_node=NUM_GPUS main_amp.py args...
+```
+`NUM_GPUS` should be less than or equal to the number of visible GPU devices on the node.  The use of `torch.distributed.launch` is unrelated to the choice of DDP wrapper.  It is safe to use either apex DDP or torch DDP with `torch.distributed.launch`.
+
+Optionally, one can run imagenet with synchronized batch normalization across processes by adding
+`--sync_bn` to the `args...`
+
+## Deterministic training (for debugging purposes)
+
+Running with the `--deterministic` flag should produce bitwise identical outputs run-to-run,
+regardless of what other options are used (see [Pytorch docs on reproducibility](https://pytorch.org/docs/stable/notes/randomness.html)).
+Since `--deterministic` disables `torch.backends.cudnn.benchmark`, `--deterministic` may
+cause a modest performance decrease.
+
+## Profiling
+
+If you're curious how the network actually looks on the CPU and GPU timelines (for example, how good is the overall utilization?
+Is the prefetcher really overlapping data transfers?) try profiling `main_amp.py`.
+[Detailed instructions can be found here](https://gist.github.com/mcarilli/213a4e698e4a0ae2234ddee56f4f3f95).
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/imagenet/main_amp.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/imagenet/main_amp.py
new file mode 100644
index 0000000000..c4b0fdfd5c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/imagenet/main_amp.py
@@ -0,0 +1,543 @@
+import argparse
+import os
+import shutil
+import time
+
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.backends.cudnn as cudnn
+import torch.distributed as dist
+import torch.optim
+import torch.utils.data
+import torch.utils.data.distributed
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+import torchvision.models as models
+
+import numpy as np
+
+try:
+    from apex.parallel import DistributedDataParallel as DDP
+    from apex.fp16_utils import *
+    from apex import amp, optimizers
+    from apex.multi_tensor_apply import multi_tensor_applier
+except ImportError:
+    raise ImportError("Please install apex from https://www.github.com/nvidia/apex to run this example.")
+
+def fast_collate(batch, memory_format):
+
+    imgs = [img[0] for img in batch]
+    targets = torch.tensor([target[1] for target in batch], dtype=torch.int64)
+    w = imgs[0].size[0]
+    h = imgs[0].size[1]
+    tensor = torch.zeros( (len(imgs), 3, h, w), dtype=torch.uint8).contiguous(memory_format=memory_format)
+    for i, img in enumerate(imgs):
+        nump_array = np.asarray(img, dtype=np.uint8)
+        if(nump_array.ndim < 3):
+            nump_array = np.expand_dims(nump_array, axis=-1)
+        nump_array = np.rollaxis(nump_array, 2)
+        tensor[i] += torch.from_numpy(nump_array)
+    return tensor, targets
+
+
+def parse():
+    model_names = sorted(name for name in models.__dict__
+                     if name.islower() and not name.startswith("__")
+                     and callable(models.__dict__[name]))
+
+    parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
+    parser.add_argument('data', metavar='DIR',
+                        help='path to dataset')
+    parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18',
+                        choices=model_names,
+                        help='model architecture: ' +
+                        ' | '.join(model_names) +
+                        ' (default: resnet18)')
+    parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
+                        help='number of data loading workers (default: 4)')
+    parser.add_argument('--epochs', default=90, type=int, metavar='N',
+                        help='number of total epochs to run')
+    parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
+                        help='manual epoch number (useful on restarts)')
+    parser.add_argument('-b', '--batch-size', default=256, type=int,
+                        metavar='N', help='mini-batch size per process (default: 256)')
+    parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
+                        metavar='LR', help='Initial learning rate.  Will be scaled by <global batch size>/256: args.lr = args.lr*float(args.batch_size*args.world_size)/256.  A warmup schedule will also be applied over the first 5 epochs.')
+    parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
+                        help='momentum')
+    parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
+                        metavar='W', help='weight decay (default: 1e-4)')
+    parser.add_argument('--print-freq', '-p', default=10, type=int,
+                        metavar='N', help='print frequency (default: 10)')
+    parser.add_argument('--resume', default='', type=str, metavar='PATH',
+                        help='path to latest checkpoint (default: none)')
+    parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
+                        help='evaluate model on validation set')
+    parser.add_argument('--pretrained', dest='pretrained', action='store_true',
+                        help='use pre-trained model')
+
+    parser.add_argument('--prof', default=-1, type=int,
+                        help='Only run 10 iterations for profiling.')
+    parser.add_argument('--deterministic', action='store_true')
+
+    parser.add_argument("--local_rank", default=os.getenv('LOCAL_RANK', 0), type=int)
+    parser.add_argument('--sync_bn', action='store_true',
+                        help='enabling apex sync BN.')
+
+    parser.add_argument('--opt-level', type=str)
+    parser.add_argument('--keep-batchnorm-fp32', type=str, default=None)
+    parser.add_argument('--loss-scale', type=str, default=None)
+    parser.add_argument('--channels-last', type=bool, default=False)
+    args = parser.parse_args()
+    return args
+
+def main():
+    global best_prec1, args
+
+    args = parse()
+    print("opt_level = {}".format(args.opt_level))
+    print("keep_batchnorm_fp32 = {}".format(args.keep_batchnorm_fp32), type(args.keep_batchnorm_fp32))
+    print("loss_scale = {}".format(args.loss_scale), type(args.loss_scale))
+
+    print("\nCUDNN VERSION: {}\n".format(torch.backends.cudnn.version()))
+
+    cudnn.benchmark = True
+    best_prec1 = 0
+    if args.deterministic:
+        cudnn.benchmark = False
+        cudnn.deterministic = True
+        torch.manual_seed(args.local_rank)
+        torch.set_printoptions(precision=10)
+
+    args.distributed = False
+    if 'WORLD_SIZE' in os.environ:
+        args.distributed = int(os.environ['WORLD_SIZE']) > 1
+
+    args.gpu = 0
+    args.world_size = 1
+
+    if args.distributed:
+        args.gpu = args.local_rank
+        torch.cuda.set_device(args.gpu)
+        torch.distributed.init_process_group(backend='nccl',
+                                             init_method='env://')
+        args.world_size = torch.distributed.get_world_size()
+
+    assert torch.backends.cudnn.enabled, "Amp requires cudnn backend to be enabled."
+
+    if args.channels_last:
+        memory_format = torch.channels_last
+    else:
+        memory_format = torch.contiguous_format
+
+    # create model
+    if args.pretrained:
+        print("=> using pre-trained model '{}'".format(args.arch))
+        model = models.__dict__[args.arch](pretrained=True)
+    else:
+        print("=> creating model '{}'".format(args.arch))
+        model = models.__dict__[args.arch]()
+
+    if args.sync_bn:
+        import apex
+        print("using apex synced BN")
+        model = apex.parallel.convert_syncbn_model(model)
+
+    model = model.cuda().to(memory_format=memory_format)
+
+    # Scale learning rate based on global batch size
+    args.lr = args.lr*float(args.batch_size*args.world_size)/256.
+    optimizer = torch.optim.SGD(model.parameters(), args.lr,
+                                momentum=args.momentum,
+                                weight_decay=args.weight_decay)
+
+    # Initialize Amp.  Amp accepts either values or strings for the optional override arguments,
+    # for convenient interoperation with argparse.
+    model, optimizer = amp.initialize(model, optimizer,
+                                      opt_level=args.opt_level,
+                                      keep_batchnorm_fp32=args.keep_batchnorm_fp32,
+                                      loss_scale=args.loss_scale
+                                      )
+
+    # For distributed training, wrap the model with apex.parallel.DistributedDataParallel.
+    # This must be done AFTER the call to amp.initialize.  If model = DDP(model) is called
+    # before model, ... = amp.initialize(model, ...), the call to amp.initialize may alter
+    # the types of model's parameters in a way that disrupts or destroys DDP's allreduce hooks.
+    if args.distributed:
+        # By default, apex.parallel.DistributedDataParallel overlaps communication with
+        # computation in the backward pass.
+        # model = DDP(model)
+        # delay_allreduce delays all communication to the end of the backward pass.
+        model = DDP(model, delay_allreduce=True)
+
+    # define loss function (criterion) and optimizer
+    criterion = nn.CrossEntropyLoss().cuda()
+
+    # Optionally resume from a checkpoint
+    if args.resume:
+        # Use a local scope to avoid dangling references
+        def resume():
+            if os.path.isfile(args.resume):
+                print("=> loading checkpoint '{}'".format(args.resume))
+                checkpoint = torch.load(args.resume, map_location = lambda storage, loc: storage.cuda(args.gpu))
+                args.start_epoch = checkpoint['epoch']
+                global best_prec1
+                best_prec1 = checkpoint['best_prec1']
+                model.load_state_dict(checkpoint['state_dict'])
+                optimizer.load_state_dict(checkpoint['optimizer'])
+                print("=> loaded checkpoint '{}' (epoch {})"
+                      .format(args.resume, checkpoint['epoch']))
+            else:
+                print("=> no checkpoint found at '{}'".format(args.resume))
+        resume()
+
+    # Data loading code
+    traindir = os.path.join(args.data, 'train')
+    valdir = os.path.join(args.data, 'val')
+
+    if(args.arch == "inception_v3"):
+        raise RuntimeError("Currently, inception_v3 is not supported by this example.")
+        # crop_size = 299
+        # val_size = 320 # I chose this value arbitrarily, we can adjust.
+    else:
+        crop_size = 224
+        val_size = 256
+
+    train_dataset = datasets.ImageFolder(
+        traindir,
+        transforms.Compose([
+            transforms.RandomResizedCrop(crop_size),
+            transforms.RandomHorizontalFlip(),
+            # transforms.ToTensor(), Too slow
+            # normalize,
+        ]))
+    val_dataset = datasets.ImageFolder(valdir, transforms.Compose([
+            transforms.Resize(val_size),
+            transforms.CenterCrop(crop_size),
+        ]))
+
+    train_sampler = None
+    val_sampler = None
+    if args.distributed:
+        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
+        val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset)
+
+    collate_fn = lambda b: fast_collate(b, memory_format)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
+        num_workers=args.workers, pin_memory=True, sampler=train_sampler, collate_fn=collate_fn)
+
+    val_loader = torch.utils.data.DataLoader(
+        val_dataset,
+        batch_size=args.batch_size, shuffle=False,
+        num_workers=args.workers, pin_memory=True,
+        sampler=val_sampler,
+        collate_fn=collate_fn)
+
+    if args.evaluate:
+        validate(val_loader, model, criterion)
+        return
+
+    for epoch in range(args.start_epoch, args.epochs):
+        if args.distributed:
+            train_sampler.set_epoch(epoch)
+
+        # train for one epoch
+        train(train_loader, model, criterion, optimizer, epoch)
+
+        # evaluate on validation set
+        prec1 = validate(val_loader, model, criterion)
+
+        # remember best prec@1 and save checkpoint
+        if args.local_rank == 0:
+            is_best = prec1 > best_prec1
+            best_prec1 = max(prec1, best_prec1)
+            save_checkpoint({
+                'epoch': epoch + 1,
+                'arch': args.arch,
+                'state_dict': model.state_dict(),
+                'best_prec1': best_prec1,
+                'optimizer' : optimizer.state_dict(),
+            }, is_best)
+
+class data_prefetcher():
+    def __init__(self, loader):
+        self.loader = iter(loader)
+        self.stream = torch.cuda.Stream()
+        self.mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255]).cuda().view(1,3,1,1)
+        self.std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255]).cuda().view(1,3,1,1)
+        # With Amp, it isn't necessary to manually convert data to half.
+        # if args.fp16:
+        #     self.mean = self.mean.half()
+        #     self.std = self.std.half()
+        self.preload()
+
+    def preload(self):
+        try:
+            self.next_input, self.next_target = next(self.loader)
+        except StopIteration:
+            self.next_input = None
+            self.next_target = None
+            return
+        # if record_stream() doesn't work, another option is to make sure device inputs are created
+        # on the main stream.
+        # self.next_input_gpu = torch.empty_like(self.next_input, device='cuda')
+        # self.next_target_gpu = torch.empty_like(self.next_target, device='cuda')
+        # Need to make sure the memory allocated for next_* is not still in use by the main stream
+        # at the time we start copying to next_*:
+        # self.stream.wait_stream(torch.cuda.current_stream())
+        with torch.cuda.stream(self.stream):
+            self.next_input = self.next_input.cuda(non_blocking=True)
+            self.next_target = self.next_target.cuda(non_blocking=True)
+            # more code for the alternative if record_stream() doesn't work:
+            # copy_ will record the use of the pinned source tensor in this side stream.
+            # self.next_input_gpu.copy_(self.next_input, non_blocking=True)
+            # self.next_target_gpu.copy_(self.next_target, non_blocking=True)
+            # self.next_input = self.next_input_gpu
+            # self.next_target = self.next_target_gpu
+
+            # With Amp, it isn't necessary to manually convert data to half.
+            # if args.fp16:
+            #     self.next_input = self.next_input.half()
+            # else:
+            self.next_input = self.next_input.float()
+            self.next_input = self.next_input.sub_(self.mean).div_(self.std)
+
+    def next(self):
+        torch.cuda.current_stream().wait_stream(self.stream)
+        input = self.next_input
+        target = self.next_target
+        if input is not None:
+            input.record_stream(torch.cuda.current_stream())
+        if target is not None:
+            target.record_stream(torch.cuda.current_stream())
+        self.preload()
+        return input, target
+
+
+def train(train_loader, model, criterion, optimizer, epoch):
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    top1 = AverageMeter()
+    top5 = AverageMeter()
+
+    # switch to train mode
+    model.train()
+    end = time.time()
+
+    prefetcher = data_prefetcher(train_loader)
+    input, target = prefetcher.next()
+    i = 0
+    while input is not None:
+        i += 1
+        if args.prof >= 0 and i == args.prof:
+            print("Profiling begun at iteration {}".format(i))
+            torch.cuda.cudart().cudaProfilerStart()
+
+        if args.prof >= 0: torch.cuda.nvtx.range_push("Body of iteration {}".format(i))
+
+        adjust_learning_rate(optimizer, epoch, i, len(train_loader))
+
+        # compute output
+        if args.prof >= 0: torch.cuda.nvtx.range_push("forward")
+        output = model(input)
+        if args.prof >= 0: torch.cuda.nvtx.range_pop()
+        loss = criterion(output, target)
+
+        # compute gradient and do SGD step
+        optimizer.zero_grad()
+
+        if args.prof >= 0: torch.cuda.nvtx.range_push("backward")
+        with amp.scale_loss(loss, optimizer) as scaled_loss:
+            scaled_loss.backward()
+        if args.prof >= 0: torch.cuda.nvtx.range_pop()
+
+        # for param in model.parameters():
+        #     print(param.data.double().sum().item(), param.grad.data.double().sum().item())
+
+        if args.prof >= 0: torch.cuda.nvtx.range_push("optimizer.step()")
+        optimizer.step()
+        if args.prof >= 0: torch.cuda.nvtx.range_pop()
+
+        if i%args.print_freq == 0:
+            # Every print_freq iterations, check the loss, accuracy, and speed.
+            # For best performance, it doesn't make sense to print these metrics every
+            # iteration, since they incur an allreduce and some host<->device syncs.
+
+            # Measure accuracy
+            prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
+
+            # Average loss and accuracy across processes for logging
+            if args.distributed:
+                reduced_loss = reduce_tensor(loss.data)
+                prec1 = reduce_tensor(prec1)
+                prec5 = reduce_tensor(prec5)
+            else:
+                reduced_loss = loss.data
+
+            # to_python_float incurs a host<->device sync
+            losses.update(to_python_float(reduced_loss), input.size(0))
+            top1.update(to_python_float(prec1), input.size(0))
+            top5.update(to_python_float(prec5), input.size(0))
+
+            torch.cuda.synchronize()
+            batch_time.update((time.time() - end)/args.print_freq)
+            end = time.time()
+
+            if args.local_rank == 0:
+                print('Epoch: [{0}][{1}/{2}]\t'
+                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                      'Speed {3:.3f} ({4:.3f})\t'
+                      'Loss {loss.val:.10f} ({loss.avg:.4f})\t'
+                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
+                      'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
+                       epoch, i, len(train_loader),
+                       args.world_size*args.batch_size/batch_time.val,
+                       args.world_size*args.batch_size/batch_time.avg,
+                       batch_time=batch_time,
+                       loss=losses, top1=top1, top5=top5))
+        if args.prof >= 0: torch.cuda.nvtx.range_push("prefetcher.next()")
+        input, target = prefetcher.next()
+        if args.prof >= 0: torch.cuda.nvtx.range_pop()
+
+        # Pop range "Body of iteration {}".format(i)
+        if args.prof >= 0: torch.cuda.nvtx.range_pop()
+
+        if args.prof >= 0 and i == args.prof + 10:
+            print("Profiling ended at iteration {}".format(i))
+            torch.cuda.cudart().cudaProfilerStop()
+            quit()
+
+
+def validate(val_loader, model, criterion):
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    top1 = AverageMeter()
+    top5 = AverageMeter()
+
+    # switch to evaluate mode
+    model.eval()
+
+    end = time.time()
+
+    prefetcher = data_prefetcher(val_loader)
+    input, target = prefetcher.next()
+    i = 0
+    while input is not None:
+        i += 1
+
+        # compute output
+        with torch.no_grad():
+            output = model(input)
+            loss = criterion(output, target)
+
+        # measure accuracy and record loss
+        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
+
+        if args.distributed:
+            reduced_loss = reduce_tensor(loss.data)
+            prec1 = reduce_tensor(prec1)
+            prec5 = reduce_tensor(prec5)
+        else:
+            reduced_loss = loss.data
+
+        losses.update(to_python_float(reduced_loss), input.size(0))
+        top1.update(to_python_float(prec1), input.size(0))
+        top5.update(to_python_float(prec5), input.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        # TODO:  Change timings to mirror train().
+        if args.local_rank == 0 and i % args.print_freq == 0:
+            print('Test: [{0}/{1}]\t'
+                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                  'Speed {2:.3f} ({3:.3f})\t'
+                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
+                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
+                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
+                   i, len(val_loader),
+                   args.world_size * args.batch_size / batch_time.val,
+                   args.world_size * args.batch_size / batch_time.avg,
+                   batch_time=batch_time, loss=losses,
+                   top1=top1, top5=top5))
+
+        input, target = prefetcher.next()
+
+    print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
+          .format(top1=top1, top5=top5))
+
+    return top1.avg
+
+
+def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
+    torch.save(state, filename)
+    if is_best:
+        shutil.copyfile(filename, 'model_best.pth.tar')
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def adjust_learning_rate(optimizer, epoch, step, len_epoch):
+    """LR schedule that should yield 76% converged accuracy with batch size 256"""
+    factor = epoch // 30
+
+    if epoch >= 80:
+        factor = factor + 1
+
+    lr = args.lr*(0.1**factor)
+
+    """Warmup"""
+    if epoch < 5:
+        lr = lr*float(1 + step + epoch*len_epoch)/(5.*len_epoch)
+
+    # if(args.local_rank == 0):
+    #     print("epoch = {}, step = {}, lr = {}".format(epoch, step, lr))
+
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the precision@k for the specified values of k"""
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].reshape(-1).float().sum(0, keepdim=True)
+        res.append(correct_k.mul_(100.0 / batch_size))
+    return res
+
+
+def reduce_tensor(tensor):
+    rt = tensor.clone()
+    dist.all_reduce(rt, op=dist.reduce_op.SUM)
+    rt /= args.world_size
+    return rt
+
+if __name__ == '__main__':
+    main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/README.md
new file mode 100644
index 0000000000..0d939cbbf6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/README.md
@@ -0,0 +1,13 @@
+**distributed_data_parallel.py** and **run.sh** show an example using Amp with
+[apex.parallel.DistributedDataParallel](https://nvidia.github.io/apex/parallel.html) or
+[torch.nn.parallel.DistributedDataParallel](https://pytorch.org/docs/stable/nn.html#distributeddataparallel)
+and the Pytorch multiprocess launcher script,
+[torch.distributed.launch](https://pytorch.org/docs/master/distributed.html#launch-utility).
+The use of `Amp` with DistributedDataParallel does not need to change from ordinary 
+single-process use.  The only gotcha is that wrapping your model with `DistributedDataParallel` must
+come after the call to `amp.initialize`.  Test via
+```bash
+bash run.sh
+```
+
+**This is intended purely as an instructional example, not a performance showcase.**
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/distributed_data_parallel.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/distributed_data_parallel.py
new file mode 100644
index 0000000000..b364405df8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/distributed_data_parallel.py
@@ -0,0 +1,65 @@
+import torch
+import argparse
+import os
+from apex import amp
+# FOR DISTRIBUTED: (can also use torch.nn.parallel.DistributedDataParallel instead)
+from apex.parallel import DistributedDataParallel
+
+parser = argparse.ArgumentParser()
+# FOR DISTRIBUTED:  Parse for the local_rank argument, which will be supplied
+# automatically by torch.distributed.launch.
+parser.add_argument("--local_rank", default=0, type=int)
+args = parser.parse_args()
+
+# FOR DISTRIBUTED:  If we are running under torch.distributed.launch,
+# the 'WORLD_SIZE' environment variable will also be set automatically.
+args.distributed = False
+if 'WORLD_SIZE' in os.environ:
+    args.distributed = int(os.environ['WORLD_SIZE']) > 1
+
+if args.distributed:
+    # FOR DISTRIBUTED:  Set the device according to local_rank.
+    torch.cuda.set_device(args.local_rank)
+
+    # FOR DISTRIBUTED:  Initialize the backend.  torch.distributed.launch will provide
+    # environment variables, and requires that you use init_method=`env://`.
+    torch.distributed.init_process_group(backend='nccl',
+                                         init_method='env://')
+
+torch.backends.cudnn.benchmark = True
+
+N, D_in, D_out = 64, 1024, 16
+
+# Each process receives its own batch of "fake input data" and "fake target data."
+# The "training loop" in each process just uses this fake batch over and over.
+# https://github.com/NVIDIA/apex/tree/master/examples/imagenet provides a more realistic
+# example of distributed data sampling for both training and validation.
+x = torch.randn(N, D_in, device='cuda')
+y = torch.randn(N, D_out, device='cuda')
+
+model = torch.nn.Linear(D_in, D_out).cuda()
+optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
+
+model, optimizer = amp.initialize(model, optimizer, opt_level="O1")
+
+if args.distributed:
+    # FOR DISTRIBUTED:  After amp.initialize, wrap the model with
+    # apex.parallel.DistributedDataParallel.
+    model = DistributedDataParallel(model)
+    # torch.nn.parallel.DistributedDataParallel is also fine, with some added args:
+    # model = torch.nn.parallel.DistributedDataParallel(model,
+    #                                                   device_ids=[args.local_rank],
+    #                                                   output_device=args.local_rank)
+
+loss_fn = torch.nn.MSELoss()
+
+for t in range(500):
+    optimizer.zero_grad()
+    y_pred = model(x)
+    loss = loss_fn(y_pred, y)
+    with amp.scale_loss(loss, optimizer) as scaled_loss:
+        scaled_loss.backward()
+    optimizer.step()
+
+if args.local_rank == 0:
+    print("final loss = ", loss)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/run.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/run.sh
new file mode 100644
index 0000000000..7a2d85f0a4
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/examples/simple/distributed/run.sh
@@ -0,0 +1,2 @@
+#!/bin/bash
+python -m torch.distributed.launch --nproc_per_node=2 distributed_data_parallel.py
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/requirements.txt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/requirements.txt
new file mode 100644
index 0000000000..f417ef5c5d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/requirements.txt
@@ -0,0 +1,5 @@
+cxxfilt>=0.2.0
+tqdm>=4.28.1
+numpy>=1.15.3
+PyYAML>=5.1
+pytest>=3.5.1
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/requirements_dev.txt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/requirements_dev.txt
new file mode 100644
index 0000000000..e1086e0e63
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/requirements_dev.txt
@@ -0,0 +1,3 @@
+-r requirements.txt
+flake8>=3.7.9
+Sphinx>=3.0.3
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/setup.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/setup.py
new file mode 100644
index 0000000000..6879ca1196
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/setup.py
@@ -0,0 +1,544 @@
+import torch
+from torch.utils.cpp_extension import BuildExtension, CppExtension, CUDAExtension, CUDA_HOME
+from setuptools import setup, find_packages
+import subprocess
+
+import sys
+import warnings
+import os
+
+# ninja build does not work unless include_dirs are abs path
+this_dir = os.path.dirname(os.path.abspath(__file__))
+
+def get_cuda_bare_metal_version(cuda_dir):
+    raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
+    output = raw_output.split()
+    release_idx = output.index("release") + 1
+    release = output[release_idx].split(".")
+    bare_metal_major = release[0]
+    bare_metal_minor = release[1][0]
+
+    return raw_output, bare_metal_major, bare_metal_minor
+
+if not torch.cuda.is_available():
+    # https://github.com/NVIDIA/apex/issues/486
+    # Extension builds after https://github.com/pytorch/pytorch/pull/23408 attempt to query torch.cuda.get_device_capability(),
+    # which will fail if you are compiling in an environment without visible GPUs (e.g. during an nvidia-docker build command).
+    print('\nWarning: Torch did not find available GPUs on this system.\n',
+          'If your intention is to cross-compile, this is not an error.\n'
+          'By default, Apex will cross-compile for Pascal (compute capabilities 6.0, 6.1, 6.2),\n'
+          'Volta (compute capability 7.0), Turing (compute capability 7.5),\n'
+          'and, if the CUDA version is >= 11.0, Ampere (compute capability 8.0).\n'
+          'If you wish to cross-compile for a single specific architecture,\n'
+          'export TORCH_CUDA_ARCH_LIST="compute capability" before running setup.py.\n')
+    if os.environ.get("TORCH_CUDA_ARCH_LIST", None) is None:
+        _, bare_metal_major, _ = get_cuda_bare_metal_version(CUDA_HOME)
+        if int(bare_metal_major) == 11:
+            os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5;8.0"
+        else:
+            os.environ["TORCH_CUDA_ARCH_LIST"] = "6.0;6.1;6.2;7.0;7.5"
+
+print("\n\ntorch.__version__  = {}\n\n".format(torch.__version__))
+TORCH_MAJOR = int(torch.__version__.split('.')[0])
+TORCH_MINOR = int(torch.__version__.split('.')[1])
+
+if TORCH_MAJOR == 0 and TORCH_MINOR < 4:
+      raise RuntimeError("Apex requires Pytorch 0.4 or newer.\n" +
+                         "The latest stable release can be obtained from https://pytorch.org/")
+
+cmdclass = {}
+ext_modules = []
+
+extras = {}
+if "--pyprof" in sys.argv:
+    string = "\n\nPyprof has been moved to its own dedicated repository and will " + \
+             "soon be removed from Apex.  Please visit\n" + \
+             "https://github.com/NVIDIA/PyProf\n" + \
+             "for the latest version."
+    warnings.warn(string, DeprecationWarning)
+    with open('requirements.txt') as f:
+        required_packages = f.read().splitlines()
+        extras['pyprof'] = required_packages
+    try:
+        sys.argv.remove("--pyprof")
+    except:
+        pass
+else:
+    warnings.warn("Option --pyprof not specified. Not installing PyProf dependencies!")
+
+if "--cpp_ext" in sys.argv or "--cuda_ext" in sys.argv:
+    if TORCH_MAJOR == 0:
+        raise RuntimeError("--cpp_ext requires Pytorch 1.0 or later, "
+                           "found torch.__version__ = {}".format(torch.__version__))
+
+if "--cpp_ext" in sys.argv:
+    sys.argv.remove("--cpp_ext")
+    ext_modules.append(
+        CppExtension('apex_C',
+                     ['csrc/flatten_unflatten.cpp',]))
+
+def get_cuda_bare_metal_version(cuda_dir):
+    raw_output = subprocess.check_output([cuda_dir + "/bin/nvcc", "-V"], universal_newlines=True)
+    output = raw_output.split()
+    release_idx = output.index("release") + 1
+    release = output[release_idx].split(".")
+    bare_metal_major = release[0]
+    bare_metal_minor = release[1][0]
+
+    return raw_output, bare_metal_major, bare_metal_minor
+
+def check_cuda_torch_binary_vs_bare_metal(cuda_dir):
+    raw_output, bare_metal_major, bare_metal_minor = get_cuda_bare_metal_version(cuda_dir)
+    torch_binary_major = torch.version.cuda.split(".")[0]
+    torch_binary_minor = torch.version.cuda.split(".")[1]
+
+    print("\nCompiling cuda extensions with")
+    print(raw_output + "from " + cuda_dir + "/bin\n")
+
+    if (bare_metal_major != torch_binary_major) or (bare_metal_minor != torch_binary_minor):
+        raise RuntimeError("Cuda extensions are being compiled with a version of Cuda that does " +
+                           "not match the version used to compile Pytorch binaries.  " +
+                           "Pytorch binaries were compiled with Cuda {}.\n".format(torch.version.cuda) +
+                           "In some cases, a minor-version mismatch will not cause later errors:  " +
+                           "https://github.com/NVIDIA/apex/pull/323#discussion_r287021798.  "
+                           "You can try commenting out this check (at your own risk).")
+
+
+# Set up macros for forward/backward compatibility hack around
+# https://github.com/pytorch/pytorch/commit/4404762d7dd955383acee92e6f06b48144a0742e
+# and
+# https://github.com/NVIDIA/apex/issues/456
+# https://github.com/pytorch/pytorch/commit/eb7b39e02f7d75c26d8a795ea8c7fd911334da7e#diff-4632522f237f1e4e728cb824300403ac
+version_ge_1_1 = []
+if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 0):
+    version_ge_1_1 = ['-DVERSION_GE_1_1']
+version_ge_1_3 = []
+if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 2):
+    version_ge_1_3 = ['-DVERSION_GE_1_3']
+version_ge_1_5 = []
+if (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4):
+    version_ge_1_5 = ['-DVERSION_GE_1_5']
+version_dependent_macros = version_ge_1_1 + version_ge_1_3 + version_ge_1_5
+
+if "--distributed_adam" in sys.argv:
+    sys.argv.remove("--distributed_adam")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--distributed_adam was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        ext_modules.append(
+            CUDAExtension(name='distributed_adam_cuda',
+                          sources=['apex/contrib/csrc/optimizers/multi_tensor_distopt_adam.cpp',
+                                   'apex/contrib/csrc/optimizers/multi_tensor_distopt_adam_kernel.cu'],
+                          include_dirs=[os.path.join(this_dir, 'csrc')],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros,
+                                              'nvcc':['-O3',
+                                                      '--use_fast_math'] + version_dependent_macros}))
+
+if "--distributed_lamb" in sys.argv:
+    sys.argv.remove("--distributed_lamb")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--distributed_lamb was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        ext_modules.append(
+            CUDAExtension(name='distributed_lamb_cuda',
+                          sources=['apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb.cpp',
+                                   'apex/contrib/csrc/optimizers/multi_tensor_distopt_lamb_kernel.cu'],
+                          include_dirs=[os.path.join(this_dir, 'csrc')],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros,
+                                              'nvcc':['-O3',
+                                                      '--use_fast_math'] + version_dependent_macros}))
+
+if "--cuda_ext" in sys.argv:
+    sys.argv.remove("--cuda_ext")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--cuda_ext was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        check_cuda_torch_binary_vs_bare_metal(CUDA_HOME)
+
+        ext_modules.append(
+            CUDAExtension(name='amp_C',
+                          sources=['csrc/amp_C_frontend.cpp',
+                                   'csrc/multi_tensor_sgd_kernel.cu',
+                                   'csrc/multi_tensor_scale_kernel.cu',
+                                   'csrc/multi_tensor_axpby_kernel.cu',
+                                   'csrc/multi_tensor_l2norm_kernel.cu',
+                                   'csrc/multi_tensor_l2norm_scale_kernel.cu',
+                                   'csrc/multi_tensor_lamb_stage_1.cu',
+                                   'csrc/multi_tensor_lamb_stage_2.cu',
+                                   'csrc/multi_tensor_adam.cu',
+                                   'csrc/multi_tensor_adagrad.cu',
+                                   'csrc/multi_tensor_novograd.cu',
+                                   'csrc/multi_tensor_lamb.cu'],
+                          extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
+                                              'nvcc':['-lineinfo',
+                                                      '-O3',
+                                                      # '--resource-usage',
+                                                      '--use_fast_math'] + version_dependent_macros}))
+        ext_modules.append(
+            CUDAExtension(name='syncbn',
+                          sources=['csrc/syncbn.cpp',
+                                   'csrc/welford.cu'],
+                          extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
+                                              'nvcc':['-O3'] + version_dependent_macros}))
+
+        ext_modules.append(
+            CUDAExtension(name='fused_layer_norm_cuda',
+                          sources=['csrc/layer_norm_cuda.cpp',
+                                   'csrc/layer_norm_cuda_kernel.cu'],
+                          extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
+                                              'nvcc':['-maxrregcount=50',
+                                                      '-O3',
+                                                      '--use_fast_math'] + version_dependent_macros}))
+
+        ext_modules.append(
+            CUDAExtension(name='mlp_cuda',
+                          sources=['csrc/mlp.cpp',
+                                   'csrc/mlp_cuda.cu'],
+                          extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
+                                              'nvcc':['-O3'] + version_dependent_macros}))
+        ext_modules.append(
+            CUDAExtension(name='fused_dense_cuda',
+                          sources=['csrc/fused_dense.cpp',
+                                   'csrc/fused_dense_cuda.cu'],
+                          extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
+                                              'nvcc':['-O3'] + version_dependent_macros}))
+
+        ext_modules.append(
+            CUDAExtension(name='scaled_upper_triang_masked_softmax_cuda',
+                          sources=['csrc/megatron/scaled_upper_triang_masked_softmax.cpp',
+                                   'csrc/megatron/scaled_upper_triang_masked_softmax_cuda.cu'],
+                          include_dirs=[os.path.join(this_dir, 'csrc')],
+                          extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
+                                              'nvcc':['-O3',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda'] + version_dependent_macros}))
+
+        ext_modules.append(
+            CUDAExtension(name='scaled_masked_softmax_cuda',
+                          sources=['csrc/megatron/scaled_masked_softmax.cpp',
+                                   'csrc/megatron/scaled_masked_softmax_cuda.cu'],
+                          include_dirs=[os.path.join(this_dir, 'csrc')],
+                          extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
+                                              'nvcc':['-O3',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda'] + version_dependent_macros}))
+
+if "--bnp" in sys.argv:
+    sys.argv.remove("--bnp")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--bnp was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        ext_modules.append(
+            CUDAExtension(name='bnp',
+                          sources=['apex/contrib/csrc/groupbn/batch_norm.cu',
+                                   'apex/contrib/csrc/groupbn/ipc.cu',
+                                   'apex/contrib/csrc/groupbn/interface.cpp',
+                                   'apex/contrib/csrc/groupbn/batch_norm_add_relu.cu'],
+                          include_dirs=[os.path.join(this_dir, 'csrc')],
+                          extra_compile_args={'cxx': [] + version_dependent_macros,
+                                              'nvcc':['-DCUDA_HAS_FP16=1',
+                                                      '-D__CUDA_NO_HALF_OPERATORS__',
+                                                      '-D__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '-D__CUDA_NO_HALF2_OPERATORS__'] + version_dependent_macros}))
+
+if "--xentropy" in sys.argv:
+    sys.argv.remove("--xentropy")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--xentropy was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        ext_modules.append(
+            CUDAExtension(name='xentropy_cuda',
+                          sources=['apex/contrib/csrc/xentropy/interface.cpp',
+                                   'apex/contrib/csrc/xentropy/xentropy_kernel.cu'],
+                          include_dirs=[os.path.join(this_dir, 'csrc')],
+                          extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
+                                              'nvcc':['-O3'] + version_dependent_macros}))
+
+if "--deprecated_fused_adam" in sys.argv:
+    sys.argv.remove("--deprecated_fused_adam")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--deprecated_fused_adam was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        ext_modules.append(
+            CUDAExtension(name='fused_adam_cuda',
+                          sources=['apex/contrib/csrc/optimizers/fused_adam_cuda.cpp',
+                                   'apex/contrib/csrc/optimizers/fused_adam_cuda_kernel.cu'],
+                          include_dirs=[os.path.join(this_dir, 'csrc')],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros,
+                                              'nvcc':['-O3',
+                                                      '--use_fast_math'] + version_dependent_macros}))
+
+if "--deprecated_fused_lamb" in sys.argv:
+    sys.argv.remove("--deprecated_fused_lamb")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--deprecated_fused_lamb was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        ext_modules.append(
+            CUDAExtension(name='fused_lamb_cuda',
+                          sources=['apex/contrib/csrc/optimizers/fused_lamb_cuda.cpp',
+                                   'apex/contrib/csrc/optimizers/fused_lamb_cuda_kernel.cu',
+                                   'csrc/multi_tensor_l2norm_kernel.cu'],
+                          include_dirs=[os.path.join(this_dir, 'csrc')],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros,
+                                              'nvcc':['-O3',
+                                                      '--use_fast_math'] + version_dependent_macros}))
+
+# Check, if ATen/CUDAGenerator.h is found, otherwise use the new ATen/CUDAGeneratorImpl.h, due to breaking change in https://github.com/pytorch/pytorch/pull/36026
+generator_flag = []
+torch_dir = torch.__path__[0]
+if os.path.exists(os.path.join(torch_dir, 'include', 'ATen', 'CUDAGenerator.h')):
+    generator_flag = ['-DOLD_GENERATOR']
+
+if "--fast_layer_norm" in sys.argv:
+    sys.argv.remove("--fast_layer_norm")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--fast_layer_norm was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        # Check, if CUDA11 is installed for compute capability 8.0
+        cc_flag = []
+        _, bare_metal_major, _ = get_cuda_bare_metal_version(CUDA_HOME)
+        if int(bare_metal_major) >= 11:
+            cc_flag.append('-gencode')
+            cc_flag.append('arch=compute_80,code=sm_80')
+
+        ext_modules.append(
+            CUDAExtension(name='fast_layer_norm',
+                          sources=['apex/contrib/csrc/layer_norm/ln_api.cpp',
+                                   'apex/contrib/csrc/layer_norm/ln_fwd_cuda_kernel.cu',
+                                   'apex/contrib/csrc/layer_norm/ln_bwd_semi_cuda_kernel.cu',
+                                   ],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros + generator_flag,
+                                              'nvcc':['-O3',
+                                                      '-gencode', 'arch=compute_70,code=sm_70',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '-U__CUDA_NO_BFLOAT16_OPERATORS__',
+                                                      '-U__CUDA_NO_BFLOAT16_CONVERSIONS__',
+                                                      '-U__CUDA_NO_BFLOAT162_OPERATORS__',
+                                                      '-U__CUDA_NO_BFLOAT162_CONVERSIONS__',
+                                                      '-I./apex/contrib/csrc/layer_norm/',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda',
+                                                      '--use_fast_math'] + version_dependent_macros + generator_flag + cc_flag},
+                          include_dirs=[os.path.join(this_dir, "apex/contrib/csrc/layer_norm")]))
+if "--fmha" in sys.argv:
+    sys.argv.remove("--fmha")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--fmha was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        # Check, if CUDA11 is installed for compute capability 8.0
+        cc_flag = []
+        _, bare_metal_major, _ = get_cuda_bare_metal_version(CUDA_HOME)
+        if int(bare_metal_major) < 11:
+            raise RuntimeError("--fmha only supported on SM80")
+
+        ext_modules.append(
+            CUDAExtension(name='fmhalib',
+                          sources=[
+                                   'apex/contrib/csrc/fmha/fmha_api.cpp',
+                                   'apex/contrib/csrc/fmha/src/fmha_noloop_reduce.cu',
+                                   'apex/contrib/csrc/fmha/src/fmha_fprop_fp16_128_64_kernel.sm80.cu',
+                                   'apex/contrib/csrc/fmha/src/fmha_fprop_fp16_256_64_kernel.sm80.cu',
+                                   'apex/contrib/csrc/fmha/src/fmha_fprop_fp16_384_64_kernel.sm80.cu',
+                                   'apex/contrib/csrc/fmha/src/fmha_fprop_fp16_512_64_kernel.sm80.cu',
+                                   'apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_128_64_kernel.sm80.cu',
+                                   'apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_256_64_kernel.sm80.cu',
+                                   'apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_384_64_kernel.sm80.cu',
+                                   'apex/contrib/csrc/fmha/src/fmha_dgrad_fp16_512_64_kernel.sm80.cu',
+                                   ],
+                          extra_compile_args={'cxx': ['-O3',
+                                                      ] + version_dependent_macros + generator_flag,
+                                              'nvcc':['-O3',
+                                                      '-gencode', 'arch=compute_80,code=sm_80',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda',
+                                                      '--use_fast_math'] + version_dependent_macros + generator_flag + cc_flag},
+                          include_dirs=[os.path.join(this_dir, "apex/contrib/csrc"), os.path.join(this_dir, "apex/contrib/csrc/fmha/src")]))
+
+
+if "--fast_multihead_attn" in sys.argv:
+    sys.argv.remove("--fast_multihead_attn")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--fast_multihead_attn was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        # Check, if CUDA11 is installed for compute capability 8.0
+        cc_flag = []
+        _, bare_metal_major, _ = get_cuda_bare_metal_version(CUDA_HOME)
+        if int(bare_metal_major) >= 11:
+            cc_flag.append('-gencode')
+            cc_flag.append('arch=compute_80,code=sm_80')
+
+        subprocess.run(["git", "submodule", "update", "--init", "apex/contrib/csrc/multihead_attn/cutlass"])
+        ext_modules.append(
+            CUDAExtension(name='fast_additive_mask_softmax_dropout',
+                          sources=['apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout.cpp',
+                                   'apex/contrib/csrc/multihead_attn/additive_masked_softmax_dropout_cuda.cu'],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros + generator_flag,
+                                              'nvcc':['-O3',
+                                                      '-gencode', 'arch=compute_70,code=sm_70',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda',
+                                                      '--use_fast_math'] + version_dependent_macros + generator_flag + cc_flag},
+                          include_dirs=[os.path.join(this_dir, "apex/contrib/csrc/multihead_attn/cutlass")]))
+        ext_modules.append(
+            CUDAExtension(name='fast_mask_softmax_dropout',
+                          sources=['apex/contrib/csrc/multihead_attn/masked_softmax_dropout.cpp',
+                                   'apex/contrib/csrc/multihead_attn/masked_softmax_dropout_cuda.cu'],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros + generator_flag,
+                                              'nvcc':['-O3',
+                                                      '-gencode', 'arch=compute_70,code=sm_70',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda',
+                                                      '--use_fast_math'] + version_dependent_macros + generator_flag + cc_flag},
+                          include_dirs=[os.path.join(this_dir, "apex/contrib/csrc/multihead_attn/cutlass")]))
+        ext_modules.append(
+            CUDAExtension(name='fast_self_multihead_attn_bias_additive_mask',
+                          sources=['apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask.cpp',
+                                   'apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_additive_mask_cuda.cu'],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros + generator_flag,
+                                              'nvcc':['-O3',
+                                                      '-gencode', 'arch=compute_70,code=sm_70',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda',
+                                                      '--use_fast_math'] + version_dependent_macros + generator_flag + cc_flag},
+                          include_dirs=[os.path.join(this_dir, "apex/contrib/csrc/multihead_attn/cutlass")]))
+        ext_modules.append(
+            CUDAExtension(name='fast_self_multihead_attn_bias',
+                          sources=['apex/contrib/csrc/multihead_attn/self_multihead_attn_bias.cpp',
+                                   'apex/contrib/csrc/multihead_attn/self_multihead_attn_bias_cuda.cu'],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros + generator_flag,
+                                              'nvcc':['-O3',
+                                                      '-gencode', 'arch=compute_70,code=sm_70',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda',
+                                                      '--use_fast_math'] + version_dependent_macros + generator_flag + cc_flag},
+                          include_dirs=[os.path.join(this_dir, "apex/contrib/csrc/multihead_attn/cutlass")]))
+        ext_modules.append(
+            CUDAExtension(name='fast_self_multihead_attn',
+                          sources=['apex/contrib/csrc/multihead_attn/self_multihead_attn.cpp',
+                                   'apex/contrib/csrc/multihead_attn/self_multihead_attn_cuda.cu'],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros + generator_flag,
+                                              'nvcc':['-O3',
+                                                      '-gencode', 'arch=compute_70,code=sm_70',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda',
+                                                      '--use_fast_math'] + version_dependent_macros + generator_flag + cc_flag},
+                          include_dirs=[os.path.join(this_dir, "apex/contrib/csrc/multihead_attn/cutlass")]))
+        ext_modules.append(
+            CUDAExtension(name='fast_self_multihead_attn_norm_add',
+                          sources=['apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add.cpp',
+                                   'apex/contrib/csrc/multihead_attn/self_multihead_attn_norm_add_cuda.cu'],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros + generator_flag,
+                                              'nvcc':['-O3',
+                                                      '-gencode', 'arch=compute_70,code=sm_70',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda',
+                                                      '--use_fast_math'] + version_dependent_macros + generator_flag + cc_flag},
+                          include_dirs=[os.path.join(this_dir, "apex/contrib/csrc/multihead_attn/cutlass")]))
+        ext_modules.append(
+            CUDAExtension(name='fast_encdec_multihead_attn',
+                          sources=['apex/contrib/csrc/multihead_attn/encdec_multihead_attn.cpp',
+                                   'apex/contrib/csrc/multihead_attn/encdec_multihead_attn_cuda.cu'],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros + generator_flag,
+                                              'nvcc':['-O3',
+                                                      '-gencode', 'arch=compute_70,code=sm_70',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda',
+                                                      '--use_fast_math'] + version_dependent_macros + generator_flag + cc_flag},
+                          include_dirs=[os.path.join(this_dir, "apex/contrib/csrc/multihead_attn/cutlass")]))
+        ext_modules.append(
+            CUDAExtension(name='fast_encdec_multihead_attn_norm_add',
+                          sources=['apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add.cpp',
+                                   'apex/contrib/csrc/multihead_attn/encdec_multihead_attn_norm_add_cuda.cu'],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros + generator_flag,
+                                              'nvcc':['-O3',
+                                                      '-gencode', 'arch=compute_70,code=sm_70',
+                                                      '-U__CUDA_NO_HALF_OPERATORS__',
+                                                      '-U__CUDA_NO_HALF_CONVERSIONS__',
+                                                      '--expt-relaxed-constexpr',
+                                                      '--expt-extended-lambda',
+                                                      '--use_fast_math'] + version_dependent_macros + generator_flag + cc_flag},
+                          include_dirs=[os.path.join(this_dir, "apex/contrib/csrc/multihead_attn/cutlass")]))
+
+if "--transducer" in sys.argv:
+    sys.argv.remove("--transducer")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--transducer was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        ext_modules.append(
+            CUDAExtension(name='transducer_joint_cuda',
+                          sources=['apex/contrib/csrc/transducer/transducer_joint.cpp',
+                                   'apex/contrib/csrc/transducer/transducer_joint_kernel.cu'],
+                          extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
+                                              'nvcc': ['-O3'] + version_dependent_macros},
+                          include_dirs=[os.path.join(this_dir, 'csrc'), os.path.join(this_dir, "apex/contrib/csrc/multihead_attn")]))
+        ext_modules.append(
+            CUDAExtension(name='transducer_loss_cuda',
+                          sources=['apex/contrib/csrc/transducer/transducer_loss.cpp',
+                                   'apex/contrib/csrc/transducer/transducer_loss_kernel.cu'],
+                          include_dirs=[os.path.join(this_dir, 'csrc')],
+                          extra_compile_args={'cxx': ['-O3'] + version_dependent_macros,
+                                              'nvcc':['-O3'] + version_dependent_macros}))
+
+if "--fast_bottleneck" in sys.argv:
+    sys.argv.remove("--fast_bottleneck")
+
+    if CUDA_HOME is None:
+        raise RuntimeError("--fast_bottleneck was requested, but nvcc was not found.  Are you sure your environment has nvcc available?  If you're installing within a container from https://hub.docker.com/r/pytorch/pytorch, only images whose names contain 'devel' will provide nvcc.")
+    else:
+        subprocess.run(["git", "submodule", "update", "--init", "apex/contrib/csrc/cudnn-frontend/"])
+        ext_modules.append(
+            CUDAExtension(name='fast_bottleneck',
+                          sources=['apex/contrib/csrc/bottleneck/bottleneck.cpp'],
+                          include_dirs=[os.path.join(this_dir, 'apex/contrib/csrc/cudnn-frontend/include')],
+                          extra_compile_args={'cxx': ['-O3',] + version_dependent_macros + generator_flag}))
+
+
+setup(
+    name='apex',
+    version='0.1',
+    packages=find_packages(exclude=('build',
+                                    'csrc',
+                                    'include',
+                                    'tests',
+                                    'dist',
+                                    'docs',
+                                    'tests',
+                                    'examples',
+                                    'apex.egg-info',)),
+    description='PyTorch Extensions written by NVIDIA',
+    ext_modules=ext_modules,
+    cmdclass={'build_ext': BuildExtension} if ext_modules else {},
+    extras_require=extras,
+)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_add_param_group.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_add_param_group.py
new file mode 100644
index 0000000000..d3e90c4330
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_add_param_group.py
@@ -0,0 +1,148 @@
+import unittest
+
+import functools as ft
+import itertools as it
+
+from apex import amp
+from apex.amp import _amp_state
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Parameter
+
+from utils import common_init, HALF, FLOAT,\
+    ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
+
+class MyModel(torch.nn.Module):
+    def __init__(self, unique):
+        super(MyModel, self).__init__()
+        self.weight0 = Parameter(unique +
+            torch.arange(2, device='cuda', dtype=torch.float32))
+        self.weight1 = Parameter(1. + unique + torch.arange(2, device='cuda', dtype=torch.float16))
+
+    @staticmethod
+    def ops(input, weight0, weight1):
+        return ((input*(weight0.float()))*(weight1.float())).sum()
+
+    def forward(self, input):
+        return self.ops(input, self.weight0, self.weight1)
+
+
+# Abandon all hope, ye who enter here.
+
+
+class TestAddParamGroup(unittest.TestCase):
+    def setUp(self):
+        self.x = torch.ones((2), device='cuda', dtype=torch.float32)
+        common_init(self)
+
+    def tearDown(self):
+        pass
+
+    def zero_grad(self, models, optimizer, how_to_zero):
+        if how_to_zero == "none":
+            for model in models:
+                for param in model.parameters():
+                    param.grad = None
+        elif how_to_zero == "model":
+            for model in models:
+                model.zero_grad()
+        elif how_to_zero == "optimizer":
+            optimizer.zero_grad()
+
+    def test_add_param_group(self):
+        for opt_level in ("O0", "O1", "O2", "O3"):
+          for zero_before_add in (True, False):
+            for try_accumulation in (True, False):
+              model0 = MyModel(1)
+              model1 = MyModel(2)
+
+              optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                          momentum=0.125)
+
+              optimizer.zero_grad()
+              loss = model0(self.x)
+              loss.backward()
+              optimizer.step()
+
+              if zero_before_add:
+                  optimizer.zero_grad()
+              optimizer.add_param_group({'params' : model1.parameters(), 'lr' : 0.5})
+              if not zero_before_add:
+                  optimizer.zero_grad()
+
+              loss = model0(self.x) + model1(self.x)
+              loss.backward(retain_graph=try_accumulation)
+              if try_accumulation:
+                  loss.backward()
+              optimizer.step()
+
+              # Once more to make sure the new params pick up momemtums properly
+              optimizer.zero_grad()
+              loss = model0(self.x) + model1(self.x)
+              loss.backward(retain_graph=try_accumulation)
+              if try_accumulation:
+                  loss.backward()
+              optimizer.step()
+
+              reference_params = [param.data.clone() for param in model0.parameters()] + \
+                                 [param.data.clone() for param in model1.parameters()]
+
+              for how_to_zero in "none", "model", "optimizer":
+                  model0 = MyModel(1)
+                  model1 = MyModel(2)
+
+                  optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                              momentum=0.125)
+
+                  _amp_state.allow_incoming_model_not_fp32 = True
+                  [model0, model1], optimizer = amp.initialize([model0, model1],
+                      optimizer,
+                      opt_level=opt_level,
+                      verbosity=0,
+                      cast_model_type=False)
+                  _amp_state.allow_incoming_model_not_fp32 = False
+
+                  _amp_state.loss_scalers[0]._loss_scale = 4.0
+
+                  self.zero_grad([model0, model1], optimizer, how_to_zero)
+                  loss = model0(self.x)
+                  with amp.scale_loss(loss, optimizer) as scaled_loss:
+                      scaled_loss.backward()
+                  optimizer.step()
+
+                  if zero_before_add:
+                      self.zero_grad([model0, model1], optimizer, how_to_zero)
+                  optimizer.add_param_group({'params' : model1.parameters(), 'lr' : 0.5})
+                  if not zero_before_add:
+                      self.zero_grad([model0, model1], optimizer, how_to_zero)
+
+                  loss = model0(self.x) + model1(self.x)
+                  with amp.scale_loss(loss, optimizer) as scaled_loss:
+                      scaled_loss.backward(retain_graph=try_accumulation)
+                  if try_accumulation:
+                      with amp.scale_loss(loss, optimizer) as scaled_loss:
+                          scaled_loss.backward()
+                  optimizer.step()
+
+                  # Once more to make sure the new params pick up momentums properly
+                  self.zero_grad([model0, model1], optimizer, how_to_zero)
+                  loss = model0(self.x) + model1(self.x)
+                  with amp.scale_loss(loss, optimizer) as scaled_loss:
+                      scaled_loss.backward(retain_graph=try_accumulation)
+                  if try_accumulation:
+                      with amp.scale_loss(loss, optimizer) as scaled_loss:
+                          scaled_loss.backward()
+                  optimizer.step()
+
+                  final_params = [param.data.clone() for param in model0.parameters()] + \
+                                 [param.data.clone() for param in model1.parameters()]
+
+                  for reference, final in zip(reference_params, final_params):
+                      self.assertTrue(torch.allclose(reference.to(final.dtype), final),
+                                      "opt_level = {}, how_to_zero = {}, zero_before_add = {}".format(
+                                      opt_level, how_to_zero, zero_before_add))
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_basic_casts.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_basic_casts.py
new file mode 100644
index 0000000000..5d4d81d1a4
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_basic_casts.py
@@ -0,0 +1,143 @@
+import unittest
+
+import functools as ft
+import itertools as it
+
+from apex import amp
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from utils import common_init, HALF, FLOAT,\
+    ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
+
+def run_layer_test(test_case, fns, expected, input_shape, test_backward=True):
+    for fn, typ in it.product(fns, expected.keys()):
+        x = torch.randn(input_shape, dtype=typ).requires_grad_()
+        y = fn(x)
+        test_case.assertEqual(y.type(), expected[typ])
+        if test_backward:
+            y.float().sum().backward()
+            test_case.assertEqual(x.grad.type(), MATCH_INPUT[typ])
+
+class TestBasicCasts(unittest.TestCase):
+    def setUp(self):
+        self.handle = amp.init(enabled=True)
+        common_init(self)
+
+    def tearDown(self):
+        self.handle._deactivate()
+
+    def test_linear_is_half(self):
+        m = nn.Linear(self.h, self.h)
+        f = ft.partial(F.linear, weight=m.weight, bias=m.bias)
+        run_layer_test(self, [m, f], ALWAYS_HALF, (self.b, self.h))
+
+    def test_conv2d_is_half(self):
+        m = nn.Conv2d(self.c, self.c, self.k)
+        f = ft.partial(F.conv2d, weight=m.weight, bias=m.bias)
+        run_layer_test(self, [m, f], ALWAYS_HALF, (self.b, self.c, self.h, self.h))
+
+    def test_softmax_is_float(self):
+        m = nn.Softmax(dim=1)
+        f = ft.partial(F.softmax, dim=1)
+        run_layer_test(self, [m, f], ALWAYS_FLOAT, (self.b, self.h))
+
+    def test_group_norm_is_float(self):
+        m = nn.GroupNorm(num_groups=4, num_channels=self.c)
+        run_layer_test(self, [m], ALWAYS_FLOAT, (self.b, self.c, self.h, self.h))
+
+    def test_mse_loss_is_float(self):
+        shape = (self.b, self.h)
+        target = torch.randn(shape)
+        mod = nn.MSELoss()
+        m = lambda x: mod(x, target)
+        f = ft.partial(F.mse_loss, target=target)
+        run_layer_test(self, [m], ALWAYS_FLOAT, shape)
+
+    def test_relu_is_match(self):
+        run_layer_test(self, [nn.ReLU(), F.relu], MATCH_INPUT, (self.b, self.h))
+
+    def test_batch_norm_is_match(self):
+        m = nn.BatchNorm2d(num_features=self.c)
+        f = ft.partial(F.batch_norm, running_mean=m.running_mean, running_var=m.running_var,
+                       weight=m.weight, bias=m.bias, training=True)
+        run_layer_test(self, [m], MATCH_INPUT, (self.b, self.c, self.h, self.h))
+
+        # Test forward-only for BN inference
+        m.eval()
+        f = ft.partial(F.batch_norm, running_mean=m.running_mean, running_var=m.running_var,
+                       weight=m.weight, bias=m.bias, training=False)
+        run_layer_test(self, [m, f], MATCH_INPUT, (self.b, self.c, self.h, self.h),
+                            test_backward=False)
+
+class TestBannedMethods(unittest.TestCase):
+    def setUp(self):
+        self.handle = amp.init(enabled=True)
+        common_init(self)
+
+    def tearDown(self):
+        self.handle._deactivate()
+
+    def bce_common(self, assertion):
+        shape = (self.b, self.h)
+        target = torch.rand(shape)
+        mod = nn.BCELoss()
+        m = lambda x: mod(x, target)
+        f = ft.partial(F.binary_cross_entropy, target=target)
+        for fn in [m, f]:
+            x = torch.rand(shape, dtype=torch.half)
+            assertion(fn, x)
+
+    def test_bce_raises_by_default(self):
+        assertion = lambda fn, x: self.assertRaises(NotImplementedError, fn, x)
+        self.bce_common(assertion)
+
+    def test_bce_is_float_with_allow_banned(self):
+        self.handle._deactivate()
+        self.handle = amp.init(enabled=True, allow_banned=True)
+        assertion = lambda fn, x: self.assertEqual(fn(x).type(), FLOAT)
+        self.bce_common(assertion)
+
+class TestTensorCasts(unittest.TestCase):
+    def setUp(self):
+        self.handle = amp.init(enabled=True)
+        common_init(self)
+
+    def tearDown(self):
+        self.handle._deactivate()
+
+    def test_matmul_method_is_half(self):
+        other = torch.randn(self.h, self.h)
+        lhs = lambda x: x.matmul(other)
+        rhs = lambda x: other.matmul(x)
+        run_layer_test(self, [lhs, rhs], ALWAYS_HALF, (self.h, self.h))
+
+    def test_matmul_op_is_half(self):
+        other = torch.randn(self.h, self.h)
+        lhs = lambda x: x @ other
+        rhs = lambda x: other @ x
+        run_layer_test(self, [lhs, rhs], ALWAYS_HALF, (self.h, self.h))
+
+    def test_pow_method_is_float(self):
+        fn = lambda x: x.pow(2.)
+        run_layer_test(self, [fn], ALWAYS_FLOAT, (self.b, self.h))
+
+    def test_pow_op_is_float(self):
+        fn = lambda x: x ** 2.
+        run_layer_test(self, [fn], ALWAYS_FLOAT, (self.b, self.h))
+
+    def test_cpu_is_float(self):
+        fn = lambda x: x.cpu()
+        always_cpu_float = {torch.float: 'torch.FloatTensor',
+                            torch.half: 'torch.FloatTensor'}
+        run_layer_test(self, [fn], always_cpu_float, (self.b, self.h))
+
+    def test_sum_is_float(self):
+        fn = lambda x: x.sum()
+        run_layer_test(self, [fn], ALWAYS_FLOAT, (self.b, self.h))
+
+    # TODO: maybe more tests on disabled casting?
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_cache.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_cache.py
new file mode 100644
index 0000000000..b58d2665f7
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_cache.py
@@ -0,0 +1,137 @@
+import unittest
+
+import functools as ft
+import itertools as it
+
+from apex import amp
+from apex.amp import _amp_state
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from utils import common_init, HALF, FLOAT,\
+    ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
+
+def get_reference_grad(i, w, ops):
+    # Creating new tensors ensures, among other things, that the new tensors are not in the cache.
+    # In fact, they are guaranteed not to use the cache because they are not torch.nn.Parameters.
+    fp32_i = i.detach().clone().float()
+    fp32_w = w.detach().clone().float().requires_grad_()
+    loss = ops(fp32_i, fp32_w)
+    loss.backward()
+    return fp32_w.grad
+
+class WhitelistModule(torch.nn.Module):
+    def __init__(self, dtype):
+        super(WhitelistModule, self).__init__()
+        self.weight = torch.nn.Parameter(torch.arange(8*8, device='cuda', dtype=dtype).view(8,8))
+
+    @staticmethod
+    def ops(input, weight):
+        return (input.mm(weight)).mm(weight).sum()
+
+    def forward(self, input):
+        return self.ops(input, self.weight)
+
+
+class BlacklistModule(torch.nn.Module):
+    def __init__(self, dtype):
+        super(BlacklistModule, self).__init__()
+        self.weight = torch.nn.Parameter(torch.arange(2*8, device='cuda', dtype=dtype).view(2,8))
+
+    @staticmethod
+    def ops(input, weight):
+        return (input + torch.pow(weight, 2) + torch.pow(weight, 2)).sum()
+
+    def forward(self, input):
+        return self.ops(input, self.weight)
+
+
+class PromoteModule(torch.nn.Module):
+    def __init__(self, dtype):
+        super(PromoteModule, self).__init__()
+        self.weight = torch.nn.Parameter(torch.arange(2*8, device='cuda', dtype=dtype).view(2,8))
+
+    @staticmethod
+    def ops(input, weight):
+        return ((input*weight)*weight).sum()
+
+    def forward(self, input):
+        return self.ops(input, self.weight)
+
+class TestCache(unittest.TestCase):
+    def setUp(self):
+        self.x = torch.ones((2, 8), device='cuda', dtype=torch.float32)
+        common_init(self)
+
+    def tearDown(self):
+        pass
+
+    def train_eval_train_test(self, module, t):
+        model = module(t).cuda()
+        optimizer = torch.optim.SGD(model.parameters(), lr=1.0)
+
+        _amp_state.allow_incoming_model_not_fp32 = True
+        model, optimizer = amp.initialize(model, optimizer, opt_level="O1", verbosity=0)
+        _amp_state.allow_incoming_model_not_fp32 = False
+        
+        def training_step():
+            for param in model.parameters():
+                param.grad = None
+        
+            loss = model(self.x).sum()
+            _amp_state.loss_scalers[0]._loss_scale = 4.0
+            with amp.scale_loss(loss, optimizer) as scaled_loss:
+                scaled_loss.backward()
+        
+            self.assertEqual(len([p.grad for p in model.parameters() if p.grad is not None]), 1)
+            self.assertEqual(model.weight.grad.type(), model.weight.type())
+        
+            reference_grad = get_reference_grad(self.x, model.weight, model.ops)
+        
+            # Currently there's no difference in the allclose calls, so no need for branching,
+            # but I'm keeping this in case we want different tolerances for fp16 and fp32 checks. 
+            if model.weight.grad.type() == "torch.cuda.HalfTensor":
+                self.assertTrue(torch.allclose(model.weight.grad.float(), reference_grad))
+            elif model.weight.grad.type() == "torch.cuda.FloatTensor":
+                self.assertTrue(torch.allclose(model.weight.grad.float(), reference_grad))
+            else:
+                raise RuntimeError("model.weight.grad.type = {}".format(model.weight.grad.type()))
+
+            model.weight.data -= 1.
+        
+        # Simulates first epoch
+        training_step()
+        
+        # Simulates eval
+        with torch.no_grad():
+            loss = model(self.x).sum()
+        
+        # Simulates resuming training after eval
+        training_step()
+
+        _amp_state.handle._deactivate()
+   
+    # I could easily have these as a set of for loops in a single test,
+    # instead of going for granularity.
+    def test_whitelist_module_fp16_weight(self):
+        self.train_eval_train_test(WhitelistModule, torch.float16)
+
+    def test_whitelist_module_fp32_weight(self):
+        self.train_eval_train_test(WhitelistModule, torch.float32)
+
+    def test_blacklist_module_fp16_weight(self):
+        self.train_eval_train_test(BlacklistModule, torch.float16)
+
+    def test_blacklist_module_fp32_weight(self):
+        self.train_eval_train_test(BlacklistModule, torch.float32)
+
+    def test_promote_module_fp16_weight(self):
+        self.train_eval_train_test(PromoteModule, torch.float16)
+
+    def test_promote_module_fp32_weight(self):
+        self.train_eval_train_test(PromoteModule, torch.float32)
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_checkpointing.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_checkpointing.py
new file mode 100644
index 0000000000..921985cd7d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_checkpointing.py
@@ -0,0 +1,267 @@
+import unittest
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+import torch.optim as optim
+
+from apex import amp
+
+
+from utils import common_init, FLOAT
+
+
+class MyModel(torch.nn.Module):
+    def __init__(self):
+        super(MyModel, self).__init__()
+        self.conv1 = nn.Conv2d(3, 6, 3, 1, 1)
+        self.bn1 = nn.BatchNorm2d(6)
+        self.param = nn.Parameter(torch.randn(1))
+
+    def forward(self, x):
+        x = x * self.param
+        x = F.relu(self.conv1(x))
+        x = self.bn1(x)
+        return x
+
+
+class TestCheckpointing(unittest.TestCase):
+    def setUp(self):
+        self.initial_lr = 1e-3
+        self.test_opt_levels = ("O0", "O1", "O2", "O3")
+
+    def seed(self):
+        torch.manual_seed(2809)
+        torch.backends.cudnn.benchmark = False
+        torch.backends.cudnn.deterministic = True
+
+    def check_state_dict_fp32(self, state_dict):
+        for key in state_dict:
+            if 'num_batches_tracked' in key:
+                continue
+            param = state_dict[key]
+            self.assertEqual(param.type(), FLOAT,
+                             'Parameter in state_dict not FLOAT')
+
+    def train_step(self, model, optimizer, data, loss_ids):
+        optimizer.zero_grad()        
+
+        output = model(data)
+
+        # Call backward for num_losses-1
+        for idx in loss_ids:
+            loss = output.mean()
+            with amp.scale_loss(loss, optimizer, loss_id=idx) as scaled_loss:
+                scaled_loss.backward(retain_graph=True)
+
+        optimizer.step()
+        return output
+
+    def compare_models(self, modelA, modelB, test_setup=''):
+        state_dictA = modelA.state_dict()
+        state_dictB = modelB.state_dict()
+        self.assertEqual(len(state_dictA), len(state_dictB),
+                         'state_dicts have different lengths' + test_setup)
+        for key in state_dictA:
+            paramA = state_dictA[key]
+            paramB = state_dictB[key]
+            self.assertTrue((paramA==paramB).all(),
+                msg='Parameters in state_dices not equal.' +
+                    'key: {}\nparam: {}\nrestored: {}\ndiff: {} for {}'.format(
+                        key, paramA, paramB, paramA - paramB, test_setup))
+
+    def test_restoring(self):
+        nb_epochs = 10
+        nb_epochs_restore = nb_epochs // 2
+        for opt_level in self.test_opt_levels:
+            for res_opt_level in self.test_opt_levels:
+                for amp_before_load in [True, False]:
+                    for num_losses in range(1, 3):
+                        test_setup = ('#' * 75 + '\n' + \
+                              f'opt_level {opt_level}\n' + \
+                              f'restore_opt_level {res_opt_level}\n' + \
+                              f'amp_before_load {amp_before_load}\n' + \
+                              f'num_losses {num_losses}\n')
+
+                        self.seed()
+
+                        # Create reference model
+                        model = MyModel().to('cuda')
+
+                        optimizer = optim.SGD(model.parameters(),
+                                              lr=self.initial_lr)
+
+                        # Initialize with num_losses*2 for the original model and the restored one
+                        model, optimizer = amp.initialize(
+                            model, optimizer, opt_level=opt_level,
+                            num_losses=num_losses*2, verbosity=0)
+
+                        # Compare training behavior for same restore option
+                        # We cannot really generalize it, since a saved model in O0
+                        # would introduce a skipped step in O1, which will raise an error
+                        if opt_level == res_opt_level:
+                            # train for nb_epochs and restore after nb_epochs_restore
+                            for epoch in range(nb_epochs):
+    
+                                x = torch.randn(16, 3, 24, 24, device='cuda')
+                                output = self.train_step(
+                                    model, optimizer, x, range(num_losses))
+                                # Initialize model one step before comparing.
+                                # Otherwise the batchnorm layers will be updated 
+                                # additionally in restore_model
+                                if epoch == (nb_epochs_restore - 1):
+                                    # Load model and optimizer
+                                    checkpoint = {
+                                        'model': model.state_dict(),
+                                        'optimizer': optimizer.state_dict(),
+                                        'amp': amp.state_dict()
+                                    }
+                                    # Check state_dict for FP32 tensors
+                                    self.check_state_dict_fp32(checkpoint['model'])
+
+                                    # Restore model
+                                    restore_model = MyModel().to('cuda')
+                                    restore_optimizer = optim.SGD(
+                                        restore_model.parameters(),
+                                        lr=self.initial_lr)
+
+                                    if amp_before_load:
+                                        restore_model, restore_optimizer = amp.initialize(
+                                            restore_model,
+                                            restore_optimizer,
+                                            opt_level=res_opt_level,
+                                            num_losses=num_losses*2,
+                                            verbosity=0)
+
+                                    restore_model.load_state_dict(checkpoint['model'])
+                                    restore_optimizer.load_state_dict(checkpoint['optimizer'])
+                                    # FIXME: We cannot test the amp.state_dict in the same script
+                                    # amp.load_state_dict(checkpoint['amp'])
+
+                                    if not amp_before_load:
+                                        restore_model, restore_optimizer = amp.initialize(
+                                            restore_model,
+                                            restore_optimizer,
+                                            opt_level=res_opt_level,
+                                            num_losses=num_losses*2,
+                                            verbosity=0)
+
+                                elif epoch >= nb_epochs_restore:
+                                    restore_output = self.train_step(
+                                        restore_model,
+                                        restore_optimizer,
+                                        x,
+                                        range(num_losses, num_losses*2))
+                                    self.assertTrue(
+                                        torch.allclose(output.float(), restore_output.float()),
+                                        'Output of reference and restored models differ for ' + test_setup)
+                                    self.compare_models(model, restore_model, test_setup)
+                        # if opt_level != res_opt_level
+                        else:
+                            # skip tests for different opt_levels
+                            continue
+
+    def test_loss_scale_decrease(self):
+        num_losses = 3
+        nb_decrease_loss_scales = [0, 1, 2]
+        for opt_level in self.test_opt_levels:
+            #print('#' * 75 + f'\n opt_level {opt_level}\n')
+            # Create new tmp copy for this run
+            nb_decrease_loss_scales_tmp = list(nb_decrease_loss_scales)
+
+            model = MyModel().to('cuda')
+        
+            optimizer = optim.SGD(model.parameters(),
+                                  lr=self.initial_lr)
+        
+            model, optimizer = amp.initialize(
+                model, optimizer, opt_level=opt_level, num_losses=num_losses,
+                verbosity=0)
+
+            if amp._amp_state.opt_properties.loss_scale != 'dynamic':
+                #print('Static loss scale set. Skipping opt_level.')
+                continue
+        
+            # force to skip some updates to decrease the loss_scale
+            initial_loss_scales = []
+            for idx in range(num_losses):
+                initial_loss_scales.append(
+                    amp._amp_state.loss_scalers[idx].loss_scale())
+            
+            for _ in range(len(nb_decrease_loss_scales)):
+                x = torch.randn(16, 3, 24, 24, device='cuda')
+                for idx in range(num_losses):
+                    while nb_decrease_loss_scales_tmp[idx] > 0:
+                        optimizer.zero_grad()
+                        output = model(x * 2**17)
+                        loss = output.mean()            
+                    
+                        with amp.scale_loss(loss, optimizer, loss_id=idx) as scaled_loss:
+                            scaled_loss.backward(retain_graph=True)
+                        optimizer.step()
+                        nb_decrease_loss_scales_tmp[idx] -= 1
+                
+            # Check loss scales afterwards
+            updated_loss_scales = []
+            for idx in range(num_losses):
+                updated_loss_scales.append(
+                    amp._amp_state.loss_scalers[idx].loss_scale())
+            for factor, update_ls, init_ls in zip(nb_decrease_loss_scales,
+                                                  updated_loss_scales,
+                                                  initial_loss_scales):
+                self.assertEqual(update_ls, init_ls / 2**factor)
+
+            # Check state dict
+            amp_state_dict = amp.state_dict()
+            for scaler_idx, factor, init_ls in zip(amp_state_dict,
+                                                   nb_decrease_loss_scales,
+                                                   initial_loss_scales):
+                scaler = amp_state_dict[scaler_idx]
+                self.assertEqual(scaler['loss_scale'], init_ls / 2**factor)
+                unskipped_target = 0
+                self.assertEqual(scaler['unskipped'], unskipped_target)
+
+    def test_state_dict(self):
+        for opt_level in self.test_opt_levels:
+            # Skip O3
+            if opt_level == 'O3':
+                continue
+
+            model = MyModel().to('cuda')
+            optimizer = optim.Adam(model.parameters(), lr=1e-3)
+            model, optimizer = amp.initialize(
+                model, optimizer, opt_level=opt_level, verbosity=0)
+
+            # Export state_dict and check for Half
+            state_dict = model.state_dict()
+            for key in state_dict:
+                self.assertFalse('Half' in state_dict[key].type())
+
+            # Check, if model is still trainable
+            # Create dummy data
+            data = torch.randn(10, 3, 4, 4, device='cuda')
+            target = torch.randn(10, 6, 4, 4, device='cuda')
+            
+            # Get initnial loss
+            optimizer.zero_grad()
+            output = model(data)
+            loss = F.mse_loss(output, target)
+            with amp.scale_loss(loss, optimizer) as scaled_loss:
+                scaled_loss.backward()
+            optimizer.step()
+            last_loss = loss.item()
+
+            # train for some epochs
+            for epoch in range(10):
+                optimizer.zero_grad()
+                output = model(data)
+                loss = F.mse_loss(output, target)
+                with amp.scale_loss(loss, optimizer) as scaled_loss:
+                    scaled_loss.backward()
+                optimizer.step()
+                self.assertTrue(loss.item() < last_loss)
+                last_loss = loss.item()
+
+if __name__=='__main__':
+    unittest.main()
+        
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_fused_sgd.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_fused_sgd.py
new file mode 100644
index 0000000000..7f592128d5
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_fused_sgd.py
@@ -0,0 +1,794 @@
+import unittest
+
+import functools as ft
+import itertools as it
+
+from apex import amp
+from apex.amp import _amp_state
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Parameter
+
+from utils import common_init, HALF, FLOAT,\
+    ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
+
+
+try:
+  import amp_C
+  disabled = False
+  from apex.optimizers import FusedSGD as FusedSGD
+except ImportError as err:
+  print("amp_C fused kernels unavailable, disabling TestMultiTensorApply.  ImportError was ", err)
+  disabled = True
+
+
+class MyModel(torch.nn.Module):
+    def __init__(self, unique):
+        super(MyModel, self).__init__()
+        self.weight0 = Parameter(unique +
+            torch.arange(2, device='cuda', dtype=torch.float32))
+        self.weight1 = Parameter(1. + unique + torch.arange(2, device='cuda', dtype=torch.float16))
+
+    @staticmethod
+    def ops(input, weight0, weight1):
+        return ((input*(weight0.float()))*(weight1.float())).sum()
+
+    def forward(self, input):
+        return self.ops(input, self.weight0, self.weight1)
+
+# Abandon all hope, ye who enter here.
+
+# This is hands down the ugliest code I have ever written, but it succeeds in testing
+# multiple models/optimizers/losses fairly thoroughly.  Many of the different test cases
+# require slightly divergent code in a way that seems near-impossible to genericize into a simple
+# cross product or nested loops.
+
+class TestMultipleModelsOptimizersLosses(unittest.TestCase):
+    def setUp(self):
+        self.x = torch.ones((2), device='cuda', dtype=torch.float32)
+        common_init(self)
+
+    def tearDown(self):
+        pass
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_2models2losses1optimizer(self):
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+
+        optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                     {'params' : model1.parameters(), 'lr' : 0.5}],
+                                    momentum=0.125)
+
+        reference_grads = []
+        for i in range(2):
+            optimizer.zero_grad()
+            loss0 = model0(self.x)
+            loss1 = model1(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads.append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()])
+
+            optimizer.step()
+
+        final_params = [param.data.clone() for param in model0.parameters()] + \
+                       [param.data.clone() for param in model1.parameters()]
+
+        for materialize_master_grads in (False, True):
+          for opt_level in ("O0", "O1", "O2", "O3"):
+            for how_to_zero in ("none", "model", "optimizer"):
+              for use_multiple_loss_scalers in (False, True):
+                if opt_level == "O1" or opt_level == "O2":
+                    inject_inf_iters = (-1, 0, 1)
+                else:
+                    inject_inf_iters = (-1,)
+
+                for inject_inf in inject_inf_iters:
+                  if inject_inf >= 0:
+                     inject_inf_locs = ("fp16", "fp32")
+                     which_backwards = (0, 1)
+                  else:
+                     inject_inf_locs = ("fdsa",)
+                     which_backwards = (None,)
+
+                  for inject_inf_loc in inject_inf_locs:
+                    for which_backward in which_backwards:
+                        if use_multiple_loss_scalers:
+                            num_losses = 2
+                            loss_ids = [0, 1]
+                        else:
+                            num_losses = 1
+                            loss_ids = [0, 0]
+
+                        if inject_inf >= 0:
+                            iters = 3
+                        else:
+                            iters = 2
+
+                        model0 = MyModel(1)
+                        model1 = MyModel(2)
+
+                        models = [model0, model1]
+
+                        optimizer = FusedSGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                              {'params' : model1.parameters(), 'lr' : 0.5}],
+                                             momentum=0.125,
+                                             materialize_master_grads=materialize_master_grads)
+
+                        _amp_state.allow_incoming_model_not_fp32 = True
+                        [model0, model1], optimizer = amp.initialize(
+                            [model0, model1],
+                            optimizer,
+                            opt_level=opt_level,
+                            verbosity=0,
+                            cast_model_type=False,
+                            num_losses=num_losses)
+                        _amp_state.allow_incoming_model_not_fp32 = False
+
+                        _amp_state.loss_scalers[0]._loss_scale = 4.0
+                        if use_multiple_loss_scalers:
+                            _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                        unskipped = 0
+                        for i in range(iters):
+                            if how_to_zero == "none":
+                                for model in models:
+                                    for param in model.parameters():
+                                        param.grad = None
+                            elif how_to_zero == "model":
+                                for model in models:
+                                    model.zero_grad()
+                            else:
+                                optimizer.zero_grad()
+
+                            loss0 = model0(self.x)
+                            loss1 = model1(self.x)
+
+                            with amp.scale_loss(loss0, optimizer, loss_id=loss_ids[0]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 0:
+                                    if inject_inf_loc == "fp32":
+                                        model0.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        model0.weight1.grad[0] = float('inf')
+                            with amp.scale_loss(loss1, optimizer, loss_id=loss_ids[1]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 1:
+                                    if inject_inf_loc == "fp32":
+                                        model1.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        model1.weight1.grad[0] = float('inf')
+
+                            if i != inject_inf:
+                                master_params = amp.master_params(optimizer)
+                                for param, reference_grad in zip(master_params, reference_grads[unskipped]):
+                                    if opt_level == "O2" and not materialize_master_grads:
+                                        continue
+                                    else:
+                                        self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()),
+                                                        "opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} use_multiple_loss_scalers {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, use_multiple_loss_scalers))
+                                unskipped += 1
+                            optimizer.step()
+
+                        model_params = [p for p in model0.parameters()] + [p for p in model1.parameters()]
+                        for model, master, reference in zip(
+                                model_params,
+                                amp.master_params(optimizer),
+                                final_params):
+                            self.assertTrue(torch.allclose(model, reference))
+                            self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                        if opt_level == "O1":
+                            _amp_state.handle._deactivate()
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_3models2losses1optimizer(self):
+
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+        model2 = MyModel(3)
+
+        optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                     {'params' : model1.parameters(), 'lr' : 0.5},
+                                     {'params' : model2.parameters(), 'lr' : 0.125}],
+                                     momentum=0.125)
+
+        reference_grads = []
+        for i in range(2):
+            optimizer.zero_grad()
+            loss0 = model0(self.x) + model2(self.x)
+            loss1 = model1(self.x) + model2(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads.append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()] +
+                                   [param.grad.data.clone() for param in model2.parameters()])
+
+            optimizer.step()
+
+
+        final_params = [param.data.clone() for param in model0.parameters()] + \
+                       [param.data.clone() for param in model1.parameters()] + \
+                       [param.data.clone() for param in model2.parameters()]
+
+        for materialize_master_grads in (False, True):
+          for opt_level in ("O0", "O1", "O2", "O3"):
+            for how_to_zero in ("none", "model", "optimizer"):
+              for use_multiple_loss_scalers in (False, True):
+                if opt_level == "O1" or opt_level == "O2":
+                    inject_inf_iters = (-1, 0, 1)
+                else:
+                    inject_inf_iters = (-1,)
+
+                for inject_inf in inject_inf_iters:
+                  if inject_inf >= 0:
+                     inject_inf_locs = ("fp16", "fp32")
+                     which_backwards = (0, 1)
+                  else:
+                     inject_inf_locs = ("fdsa",)
+                     which_backwards = (None,)
+
+                  for inject_inf_loc in inject_inf_locs:
+                    for which_backward in which_backwards:
+                      if use_multiple_loss_scalers:
+                          num_losses = 2
+                          loss_ids = [0, 1]
+                      else:
+                          num_losses = 1
+                          loss_ids = [0, 0]
+
+                      if inject_inf >= 0:
+                          iters = 3
+                          if which_backward == 0:
+                              which_models = (0, 2)
+                          elif which_backward == 1:
+                              which_models = (1, 2)
+                      else:
+                          iters = 2
+                          which_models = (None,)
+
+                      for which_model in which_models:
+                          model0 = MyModel(1)
+                          model1 = MyModel(2)
+                          model2 = MyModel(3)
+
+                          models = [model0, model1, model2]
+
+                          optimizer = FusedSGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                               {'params' : model1.parameters(), 'lr' : 0.5},
+                                               {'params' : model2.parameters(), 'lr' : 0.125}],
+                                               momentum=0.125,
+                                               materialize_master_grads=materialize_master_grads)
+
+                          _amp_state.allow_incoming_model_not_fp32 = True
+                          [model0, model1, model2], optimizer = amp.initialize(
+                              [model0, model1, model2],
+                              optimizer,
+                              opt_level=opt_level,
+                              verbosity=0,
+                              cast_model_type=False,
+                              num_losses=num_losses)
+                          _amp_state.allow_incoming_model_not_fp32 = False
+
+                          _amp_state.loss_scalers[0]._loss_scale = 4.0
+                          if use_multiple_loss_scalers:
+                              _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                          unskipped = 0
+                          for i in range(iters):
+                              if how_to_zero == "none":
+                                  for model in models:
+                                      for param in model.parameters():
+                                          param.grad = None
+                              elif how_to_zero == "model":
+                                  for model in models:
+                                      model.zero_grad()
+                              else:
+                                  optimizer.zero_grad()
+
+                              loss0 = model0(self.x) + model2(self.x)
+                              loss1 = model1(self.x) + model2(self.x)
+
+                              with amp.scale_loss(loss0, optimizer, loss_id=loss_ids[0]) as scaled_loss:
+                                  scaled_loss.backward()
+                                  if i == inject_inf and which_backward == 0:
+                                      if which_model == 0:
+                                          inj_model = model0
+                                      elif which_model == 2:
+                                          inj_model = model2
+                                      else:
+                                          raise RuntimeError(which_model + " invalid for loss 0")
+                                      if inject_inf_loc == "fp32":
+                                          inj_model.weight0.grad[0] = float('inf')
+                                      elif inject_inf_loc == "fp16":
+                                          inj_model.weight1.grad[0] = float('inf')
+                              with amp.scale_loss(loss1, optimizer, loss_id=loss_ids[1]) as scaled_loss:
+                                  scaled_loss.backward()
+                                  if i == inject_inf and which_backward == 1:
+                                      if which_model == 1:
+                                          inj_model = model1
+                                      elif which_model == 2:
+                                          inj_model = model2
+                                      else:
+                                          raise RuntimeError(which_model + " invalid for loss 1 ")
+                                      if inject_inf_loc == "fp32":
+                                          inj_model.weight0.grad[0] = float('inf')
+                                      elif inject_inf_loc == "fp16":
+                                          inj_model.weight1.grad[0] = float('inf')
+
+                              if i != inject_inf:
+                                  master_params = amp.master_params(optimizer)
+                                  for param, reference_grad in zip(master_params, reference_grads[unskipped]):
+                                      if opt_level == "O2" and not materialize_master_grads:
+                                          continue
+                                      else:
+                                          self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()),
+                                          "opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} which_model {} use_multiple_loss_scalers {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, which_model, use_multiple_loss_scalers))
+                                  unskipped += 1
+
+                              optimizer.step()
+
+                          model_params = [p for p in model0.parameters()] + \
+                                         [p for p in model1.parameters()] + \
+                                         [p for p in model2.parameters()]
+                          for model, master, reference in zip(
+                                  model_params,
+                                  amp.master_params(optimizer),
+                                  final_params):
+                              self.assertTrue(torch.allclose(model, reference))
+                              self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                          if opt_level == "O1":
+                              _amp_state.handle._deactivate()
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_2models2losses2optimizers(self):
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+
+        optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                      momentum=0.125)
+        optimizer1 = torch.optim.SGD([{'params' : model1.parameters(), 'lr' : 0.5}],
+                                      momentum=0.25)
+
+        # Don't do it like this:  reference_grads = [[]]*5
+        # because then it creates a list of 5 references to the same "[]" and appending
+        # to any of them effectively makes you append to all of them, which multiplies
+        # the resulting size of reference_grads by 5x and needless to say makes the test fail.
+        reference_grads = [[], [], [], [], []]
+        final_params = [None, None, None, None, None]
+        for i in range(2):
+            optimizer0.zero_grad()
+            optimizer1.zero_grad()
+            loss0 = model0(self.x)
+            loss1 = model1(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads[0].append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()])
+
+            optimizer0.step()
+            optimizer1.step()
+
+        final_params[0] = [param.data.clone() for param in model0.parameters()] + \
+                          [param.data.clone() for param in model1.parameters()]
+
+        def what_got_skipped(which_iter, which_backward):
+            if which_iter == 0 and which_backward == 0:
+                return 1
+            if which_iter == 0 and which_backward == 1:
+                return 2
+            if which_iter == 1 and which_backward == 0:
+                return 3
+            if which_iter == 1 and which_backward == 1:
+                return 4
+            return 0
+
+        for which_iter in (0,1):
+            for which_backward in (0,1):
+                model0 = MyModel(1)
+                model1 = MyModel(2)
+
+                optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                              momentum=0.125)
+                optimizer1 = torch.optim.SGD([{'params' : model1.parameters(), 'lr' : 0.5}],
+                                              momentum=0.25)
+
+                for i in range(3):
+                    optimizer0.zero_grad()
+                    optimizer1.zero_grad()
+                    loss0 = model0(self.x)
+                    loss1 = model1(self.x)
+                    loss0.backward()
+                    loss1.backward()
+
+                    if i != which_iter:
+                        reference_grads[what_got_skipped(which_iter, which_backward)].append(
+                            [param.grad.data.clone() for param in model0.parameters()] +
+                            [param.grad.data.clone() for param in model1.parameters()])
+
+                    if i == which_iter:
+                        if which_backward == 0:
+                            optimizer1.step()
+                        else:
+                            optimizer0.step()
+                    else:
+                        optimizer0.step()
+                        optimizer1.step()
+
+                final_params[what_got_skipped(which_iter, which_backward)] = \
+                    [param.data.clone() for param in model0.parameters()] + \
+                    [param.data.clone() for param in model1.parameters()]
+
+        for materialize_master_grads in (False, True):
+          for opt_level in ("O0", "O1", "O2", "O3"):
+            for how_to_zero in ("none", "model", "optimizer"):
+              for use_multiple_loss_scalers in (False, True):
+                if opt_level == "O1" or opt_level == "O2":
+                    inject_inf_iters = (-1, 0, 1)
+                else:
+                    inject_inf_iters = (-1,)
+
+                for inject_inf in inject_inf_iters:
+                  if inject_inf >= 0:
+                     inject_inf_locs = ("fp16", "fp32")
+                     which_backwards = (0, 1)
+                  else:
+                     inject_inf_locs = ("fdsa",)
+                     which_backwards = (None,)
+
+                  for inject_inf_loc in inject_inf_locs:
+                    for which_backward in which_backwards:
+                        if use_multiple_loss_scalers:
+                            num_losses = 2
+                            loss_ids = [0, 1]
+                        else:
+                            num_losses = 1
+                            loss_ids = [0, 0]
+
+                        if inject_inf >= 0:
+                            iters = 3
+                        else:
+                            iters = 2
+
+                        model0 = MyModel(1)
+                        model1 = MyModel(2)
+
+                        models = [model0, model1]
+
+                        optimizer0 = FusedSGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                              momentum=0.125, materialize_master_grads=materialize_master_grads)
+                        optimizer1 = FusedSGD([{'params' : model1.parameters(), 'lr' : 0.5}],
+                                              momentum=0.25, materialize_master_grads=materialize_master_grads)
+
+                        _amp_state.allow_incoming_model_not_fp32 = True
+                        [model0, model1], [optimizer0, optimizer1] = amp.initialize(
+                            [model0, model1],
+                            [optimizer0, optimizer1],
+                            opt_level=opt_level,
+                            verbosity=0,
+                            cast_model_type=False,
+                            num_losses=num_losses)
+                        _amp_state.allow_incoming_model_not_fp32 = False
+
+                        _amp_state.loss_scalers[0]._loss_scale = 4.0
+                        if use_multiple_loss_scalers:
+                            _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                        unskipped = 0
+                        for i in range(iters):
+                            if how_to_zero == "none":
+                                for model in models:
+                                    for param in model.parameters():
+                                        param.grad = None
+                            elif how_to_zero == "model":
+                                for model in models:
+                                    model.zero_grad()
+                            else:
+                                optimizer0.zero_grad()
+                                optimizer1.zero_grad()
+
+                            loss0 = model0(self.x)
+                            loss1 = model1(self.x)
+
+                            with amp.scale_loss(loss0, optimizer0, loss_id=loss_ids[0]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 0:
+                                    if inject_inf_loc == "fp32":
+                                        model0.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        model0.weight1.grad[0] = float('inf')
+                            with amp.scale_loss(loss1, optimizer1, loss_id=loss_ids[1]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 1:
+                                    if inject_inf_loc == "fp32":
+                                        model1.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        model1.weight1.grad[0] = float('inf')
+
+                            # print("opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} use_multiple_loss_scalers {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, use_multiple_loss_scalers))
+
+                            if i != inject_inf:
+                                master_params = list(amp.master_params(optimizer0)) + \
+                                                list(amp.master_params(optimizer1))
+                                for param, reference_grad in zip(master_params,
+                                        reference_grads[what_got_skipped(inject_inf, which_backward)][unskipped]):
+                                    if opt_level == "O2" and not materialize_master_grads:
+                                        continue
+                                    else:
+                                        self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()))
+                                unskipped += 1
+
+                            optimizer0.step()
+                            optimizer1.step()
+
+                        model_params = [p for p in model0.parameters()] + [p for p in model1.parameters()]
+                        master_params = [p for p in amp.master_params(optimizer0)] + \
+                                        [p for p in amp.master_params(optimizer1)]
+                        for model, master, reference in zip(
+                                model_params,
+                                master_params,
+                                final_params[what_got_skipped(inject_inf, which_backward)]):
+                            self.assertTrue(torch.allclose(model, reference))
+                            self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                        if opt_level == "O1":
+                            _amp_state.handle._deactivate()
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_3models2losses2optimizers(self):
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+        model2 = MyModel(3)
+
+        optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                      {'params' : model1.parameters(), 'lr' : 1.0}],
+                                     momentum=0.5)
+        optimizer1 = torch.optim.SGD([{'params' : model2.parameters(), 'lr' : 0.5}],
+                                     momentum=0.25)
+
+        # Again, can't do this:  reference_grads = [[]]*9
+        reference_grads = [[], [], [], [], [], [], [], [], []]
+        final_params = [None, None, None, None, None, None, None, None, None]
+        for i in range(2):
+            optimizer0.zero_grad()
+            optimizer1.zero_grad()
+            loss0 = model0(self.x) + model1(self.x)
+            loss1 = model2(self.x) + model1(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads[0].append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()])
+
+            optimizer0.step()
+            optimizer1.step()
+
+        final_params[0] = \
+            [param.data.clone() for param in model0.parameters()] + \
+            [param.data.clone() for param in model1.parameters()] + \
+            [param.data.clone() for param in model2.parameters()]
+
+        def what_got_skipped(which_iter, which_backward, which_model):
+            if which_iter == 0:
+                if which_backward == 0:
+                    if which_model == 0:
+                        return 1
+                    if which_model == 1:
+                        return 2
+                if which_backward == 1:
+                    if which_model == 2:
+                        return 3
+                    if which_model == 1:
+                        return 4
+            if which_iter == 1:
+                if which_backward == 0:
+                    if which_model == 0:
+                        return 5
+                    if which_model == 1:
+                        return 6
+                if which_backward == 1:
+                    if which_model == 2:
+                        return 7
+                    if which_model == 1:
+                        return 8
+            return 0
+
+        for which_iter in (0,1):
+            for which_backward in (0,1):
+                if which_backward == 0:
+                    which_models = (0,1)
+                if which_backward == 1:
+                    which_models = (2,1)
+                for which_model in which_models:
+
+                    model0 = MyModel(1)
+                    model1 = MyModel(2)
+                    model2 = MyModel(3)
+
+                    optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                                  {'params' : model1.parameters(), 'lr' : 1.0}],
+                                                 momentum=0.5)
+                    optimizer1 = torch.optim.SGD([{'params' : model2.parameters(), 'lr' : 0.5}],
+                                                 momentum=0.25)
+
+                    for i in range(3):
+                        optimizer0.zero_grad()
+                        optimizer1.zero_grad()
+                        loss0 = model0(self.x) + model1(self.x)
+                        loss1 = model2(self.x) + model1(self.x)
+                        loss0.backward()
+                        loss1.backward()
+
+                        if i != which_iter:
+                            reference_grads[what_got_skipped(which_iter,
+                                    which_backward, which_model)].append(
+                                [param.grad.data.clone() for param in model0.parameters()] +
+                                [param.grad.data.clone() for param in model1.parameters()])
+
+                        if i == which_iter:
+                            if which_backward == 0:
+                                # if which_model == 0:
+                                    optimizer1.step()
+                                # if which_model == 1:
+                                #     optimizer1.step()
+                            if which_backward == 1:
+                                # if which_model == 2:
+                                #     optimizer0.step()
+                                # if which_model == 1:
+                                    continue
+                        else:
+                            optimizer0.step()
+                            optimizer1.step()
+
+                    final_params[what_got_skipped(which_iter, which_backward, which_model)] = \
+                        [param.data.clone() for param in model0.parameters()] + \
+                        [param.data.clone() for param in model1.parameters()] + \
+                        [param.data.clone() for param in model2.parameters()]
+
+        for materialize_master_grads in (False, True):
+          for opt_level in ("O0", "O1", "O2", "O3"):
+            for how_to_zero in ("none", "model", "optimizer"):
+              for use_multiple_loss_scalers in (False, True):
+                if opt_level == "O1" or opt_level == "O2":
+                    inject_inf_iters = (-1, 0, 1)
+                else:
+                    inject_inf_iters = (-1,)
+
+                for inject_inf in inject_inf_iters:
+                  if inject_inf >= 0:
+                     inject_inf_locs = ("fp16", "fp32")
+                     which_backwards = (0, 1)
+                  else:
+                     inject_inf_locs = ("fdsa",)
+                     which_backwards = (None,)
+
+                  for inject_inf_loc in inject_inf_locs:
+                    for which_backward in which_backwards:
+                      if use_multiple_loss_scalers:
+                          num_losses = 2
+                          loss_ids = [0, 1]
+                      else:
+                          num_losses = 1
+                          loss_ids = [0, 0]
+
+                      if inject_inf >= 0:
+                          iters = 3
+                          if which_backward == 0:
+                              which_models = (0, 1)
+                          elif which_backward == 1:
+                              which_models = (2, 1)
+                      else:
+                          iters = 2
+                          which_models = (None,)
+
+                      for which_model in which_models:
+                          model0 = MyModel(1)
+                          model1 = MyModel(2)
+                          model2 = MyModel(3)
+
+                          models = [model0, model1, model2]
+
+                          optimizer0 = FusedSGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                            {'params' : model1.parameters(), 'lr' : 1.0}],
+                                            momentum=0.5, materialize_master_grads=materialize_master_grads)
+                          optimizer1 = FusedSGD([{'params' : model2.parameters(), 'lr' : 0.5}],
+                                                momentum=0.25, materialize_master_grads=materialize_master_grads)
+
+                          _amp_state.allow_incoming_model_not_fp32 = True
+                          [model0, model1, model2], [optimizer0, optimizer1] = amp.initialize(
+                              [model0, model1, model2],
+                              [optimizer0, optimizer1],
+                              opt_level=opt_level,
+                              verbosity=0,
+                              cast_model_type=False,
+                              num_losses=num_losses)
+                          _amp_state.allow_incoming_model_not_fp32 = False
+
+                          _amp_state.loss_scalers[0]._loss_scale = 4.0
+                          if use_multiple_loss_scalers:
+                              _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                          unskipped = 0
+                          for i in range(iters):
+                              if how_to_zero == "none":
+                                  for model in models:
+                                      for param in model.parameters():
+                                          param.grad = None
+                              elif how_to_zero == "model":
+                                  for model in models:
+                                      model.zero_grad()
+                              else:
+                                  optimizer0.zero_grad()
+                                  optimizer1.zero_grad()
+
+                              loss0 = model0(self.x) + model1(self.x)
+                              loss1 = model2(self.x) + model1(self.x)
+
+                              with amp.scale_loss(loss0, optimizer0, loss_id=loss_ids[0]) as scaled_loss:
+                                  scaled_loss.backward()
+                                  if i == inject_inf and which_backward == 0:
+                                      if which_model == 0:
+                                          inj_model = model0
+                                      elif which_model == 1:
+                                          inj_model = model1
+                                      else:
+                                          raise RuntimeError(which_model + " invalid for loss 0")
+                                      if inject_inf_loc == "fp32":
+                                          inj_model.weight0.grad[0] = float('inf')
+                                      elif inject_inf_loc == "fp16":
+                                          inj_model.weight1.grad[0] = float('inf')
+                              with amp.scale_loss(loss1, [optimizer0, optimizer1], loss_id=loss_ids[1]) as scaled_loss:
+                                  scaled_loss.backward()
+                                  if i == inject_inf and which_backward == 1:
+                                      if which_model == 2:
+                                          inj_model = model2
+                                      elif which_model == 1:
+                                          inj_model = model1
+                                      else:
+                                          raise RuntimeError(which_model + " invalid for loss 1 ")
+                                      if inject_inf_loc == "fp32":
+                                          inj_model.weight0.grad[0] = float('inf')
+                                      elif inject_inf_loc == "fp16":
+                                          inj_model.weight1.grad[0] = float('inf')
+
+                              if i != inject_inf:
+                                  master_params = list(amp.master_params(optimizer0)) + \
+                                                  list(amp.master_params(optimizer1))
+                                  for param, reference_grad in zip(master_params,
+                                        reference_grads[what_got_skipped(inject_inf,
+                                            which_backward, which_model)][unskipped]):
+                                      if opt_level == "O2" and not materialize_master_grads:
+                                          continue
+                                      else:
+                                          self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()))
+                                  unskipped += 1
+
+                              optimizer0.step()
+                              optimizer1.step()
+
+                          model_params = [p for p in model0.parameters()] + \
+                                         [p for p in model1.parameters()] + \
+                                         [p for p in model2.parameters()]
+                          master_params = [p for p in amp.master_params(optimizer0)] + \
+                                          [p for p in amp.master_params(optimizer1)]
+
+                          # print("opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} use_multiple_loss_scalers {} which_model {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, use_multiple_loss_scalers, which_model))
+
+                          for model, master, reference in zip(
+                                  model_params,
+                                  master_params,
+                                  final_params[what_got_skipped(inject_inf, which_backward, which_model)]):
+                              self.assertTrue(torch.allclose(model, reference))
+                              self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                          if opt_level == "O1":
+                              _amp_state.handle._deactivate()
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_larc.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_larc.py
new file mode 100644
index 0000000000..f4f3e838f6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_larc.py
@@ -0,0 +1,53 @@
+import unittest
+
+import torch
+from torch import nn
+from torch.nn import Parameter
+
+from apex import amp
+from apex.parallel.LARC import LARC
+from utils import common_init
+
+
+class MyModel(torch.nn.Module):
+    def __init__(self, unique):
+        super(MyModel, self).__init__()
+        self.weight0 = Parameter(
+            unique + torch.arange(2, device="cuda", dtype=torch.float32)
+        )
+
+    def forward(self, input):
+        return (input * self.weight0).sum()
+
+
+class TestLARC(unittest.TestCase):
+    def setUp(self):
+        self.x = torch.ones((2), device="cuda", dtype=torch.float32)
+        common_init(self)
+
+    def tearDown(self):
+        pass
+
+    def test_larc_mixed_precision(self):
+        for opt_level in ["O0", "O1", "O2", "O3"]:
+            model = MyModel(1)
+
+            optimizer = LARC(
+                torch.optim.SGD(
+                    [{"params": model.parameters(), "lr": 0.25}], momentum=0.125
+                )
+            )
+
+            model, optimizer = amp.initialize(
+                model, optimizer, opt_level=opt_level, verbosity=0
+            )
+
+            optimizer.zero_grad()
+            loss = model(self.x)
+            with amp.scale_loss(loss, optimizer) as scaled_loss:
+                scaled_loss.backward()
+            optimizer.step()
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_axpby.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_axpby.py
new file mode 100644
index 0000000000..0b439bb8d9
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_axpby.py
@@ -0,0 +1,180 @@
+import unittest
+
+import functools as ft
+import itertools as it
+
+from apex import amp
+import torch
+from torch import nn
+import torch.nn.functional as F
+from math import floor
+
+from utils import common_init, HALF, FLOAT,\
+    ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
+
+try:
+  import amp_C
+  from amp_C import multi_tensor_axpby
+  from apex.multi_tensor_apply import MultiTensorApply
+  disabled = False
+except ImportError as err:
+  print("amp_C fused kernels unavailable, disabling TestMultiTensorApply.  ImportError was ", err)
+  disabled = True
+
+TORCH_MAJOR = int(torch.__version__.split('.')[0])
+TORCH_MINOR = int(torch.__version__.split('.')[1])
+try_nhwc = (TORCH_MAJOR > 1) or (TORCH_MAJOR == 1 and TORCH_MINOR > 4)
+
+
+class TestMultiTensorAxpby(unittest.TestCase):
+
+    def setUp(self):
+        common_init(self)
+
+        self.a = 2.0
+        self.b = 8.0
+        self.xval = 4.0
+        self.yval = 16.0
+        self.overflow_buf = torch.cuda.IntTensor(1).zero_()
+        self.ref = torch.full((1,), 136.0, device="cuda", dtype=torch.float32)
+
+    def tearDown(self):
+        pass
+
+    # The tensor creation here is written for convenience, not speed.
+    def axpby(self, sizea, sizeb, applier, repeat_tensors,
+              x_type, y_type, out_type, inplace=False, nhwc=False):
+        self.overflow_buf.zero_()
+        sizea = sizea if isinstance(sizea, tuple) else (sizea,)
+        sizeb = sizeb if isinstance(sizeb, tuple) else (sizeb,)
+        t1 = torch.full(sizea, 1.0, device="cuda", dtype=torch.float32)
+        t2 = torch.full(sizeb, 1.0, device="cuda", dtype=torch.float32)
+
+        def to_fmt(t, tp):
+            if nhwc:
+                return t.clone().to(tp, memory_format=torch.channels_last)
+            else:
+                return t.clone().to(tp)
+
+        y_list = []
+        for i in range(repeat_tensors):
+            y_list += [to_fmt(t1, y_type)*self.yval, to_fmt(t2, y_type)*self.yval]
+
+        x_list = [to_fmt(x, x_type)*(self.xval/self.yval) for x in y_list]
+
+        if inplace:
+            out_list = y_list
+        else:
+            out_list = [to_fmt(out, out_type)*3.0 for out in y_list]
+
+        applier(multi_tensor_axpby, self.overflow_buf, [x_list, y_list, out_list], self.a, self.b, -1)
+
+        self.assertTrue(all([torch.allclose(out, self.ref.to(out_type)) for out in out_list]),
+                        msg="{} {} {} {} {} {} {}".format(sizea, sizeb, repeat_tensors,
+                        x_type, y_type, out_type, inplace))
+        self.assertTrue(self.overflow_buf.item() == 0,
+                        msg="{} {} {} {} {} {} {}".format(sizea, sizeb, repeat_tensors,
+                        x_type, y_type, out_type, inplace))
+
+    # def find_inf(self, sizea, sizeb, applier, repeat_tensors, in_type, out_type, t, ind, val, inplace=False):
+    #     self.overflow_buf.zero_()
+    #     a = torch.cuda.FloatTensor(sizea).fill_(self.scale)
+    #     b = torch.cuda.FloatTensor(sizeb).fill_(self.scale)
+
+    #     out_list = []
+    #     for i in range(repeat_tensors):
+    #         out_list += [a.clone().to(out_type), b.clone().to(out_type)]
+
+    #     if inplace:
+    #         in_list = out_list
+    #     else:
+    #         in_list = [out.clone().to(in_type) for out in out_list]
+
+    #     applier(multi_tensor_scale, self.overflow_buf, [in_list, out_list], 1./self.scale)
+
+    #     self.overflow_buf.zero_()
+    #     in_list[t][ind] = val
+    #     applier(multi_tensor_scale, self.overflow_buf, [in_list, out_list], 1./self.scale)
+    #     self.assertTrue(self.overflow_buf.item())
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_fuzz(self):
+        input_size_pairs = (
+            (7777*77, 555*555),
+            (777, 555),
+            (555, 2048*32+1),
+            (2048*32+1, 555),
+            (555, 2048*32),
+            (2048*32, 555),
+            (33333, 555),
+            (555, 33333))
+        appliers = (
+            MultiTensorApply(2048*32),
+            MultiTensorApply(333),
+            MultiTensorApply(33333))
+        repeat_tensors = (
+            1,
+            55)
+
+        for sizea, sizeb in input_size_pairs:
+          for applier in appliers:
+            for repeat in repeat_tensors:
+              for x_type in (torch.float32, torch.float16):
+                for y_type in (torch.float32, torch.float16):
+                  for out_type in (torch.float32, torch.float16):
+                    for inplace in (True, False):
+                      if inplace is True and (y_type is not out_type):
+                        continue
+                      else:
+                        self.axpby(sizea, sizeb, applier, repeat,
+                                   x_type, y_type, out_type, inplace=inplace)
+                      # self.find_inf(sizea, sizeb, applier, repeat, in_type, out_type,
+                      #               0, 0, float('nan'), inplace=inplace)
+                      # self.find_inf(sizea, sizeb, applier, repeat, in_type, out_type,
+                      #               2*repeat-1, sizeb-1, float('inf'), inplace=inplace)
+                      # self.find_inf(sizea, sizeb, applier, repeat, in_type, out_type,
+                      #              2*(repeat//2), sizea//2, float('inf'), inplace=inplace)
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    @unittest.skipIf(not try_nhwc, "torch version is 1.4 or earlier, may not support nhwc")
+    def test_fuzz_nhwc(self):
+        input_size_pairs = (
+            ((7, 77, 7, 77), (5, 55, 5, 55)),
+            ((1, 1, 777, 1), (1, 1, 555, 1)),
+            ((5, 47, 5, 55), (1, 1, 1, 2048*32 + 1)),
+            ((1, 1, 1, 2048*32 + 1), (55, 47, 5, 55)),
+            ((555, 1, 1, 1), (32, 8, 32, 8)),
+            ((32, 8, 32, 8), (55, 47, 5, 55)),
+            ((1, 1, 33333, 1), (55, 47, 55, 5)),
+            ((55, 47, 55, 5), (1, 1, 33333, 1)))
+        appliers = (
+            MultiTensorApply(2048*32),
+            MultiTensorApply(333),
+            MultiTensorApply(33333))
+        repeat_tensors = (
+            1,
+            55)
+
+        for sizea, sizeb in input_size_pairs:
+          for applier in appliers:
+            for repeat in repeat_tensors:
+              for x_type in (torch.float32, torch.float16):
+                for y_type in (torch.float32, torch.float16):
+                  for out_type in (torch.float32, torch.float16):
+                    for inplace in (True, False):
+                      if inplace is True and (y_type is not out_type):
+                        continue
+                      else:
+                        self.axpby(sizea, sizeb, applier, repeat,
+                                   x_type, y_type, out_type, inplace=inplace, nhwc=True)
+                      # self.find_inf(sizea, sizeb, applier, repeat, in_type, out_type,
+                      #               0, 0, float('nan'), inplace=inplace)
+                      # self.find_inf(sizea, sizeb, applier, repeat, in_type, out_type,
+                      #               2*repeat-1, sizeb-1, float('inf'), inplace=inplace)
+                      # self.find_inf(sizea, sizeb, applier, repeat, in_type, out_type,
+                      #              2*(repeat//2), sizea//2, float('inf'), inplace=inplace)
+
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_l2norm.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_l2norm.py
new file mode 100644
index 0000000000..ed3cbd1956
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_l2norm.py
@@ -0,0 +1,87 @@
+import unittest
+
+import functools as ft
+import itertools as it
+
+from apex import amp
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from utils import common_init, HALF, FLOAT,\
+    ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
+
+try:
+  import amp_C
+  from amp_C import multi_tensor_l2norm
+  from apex.multi_tensor_apply import MultiTensorApply
+  disabled = False
+except ImportError as err:
+  print("amp_C fused kernels unavailable, disabling TestMultiTensorApply.  ImportError was ", err)
+  disabled = True
+
+
+class TestMultiTensorL2Norm(unittest.TestCase):
+
+    def setUp(self):
+        common_init(self)
+        self.val = 4.0
+        self.overflow_buf = torch.cuda.IntTensor(1).zero_()
+
+    def tearDown(self):
+        pass
+
+    # The tensor creation here is written for convenience, not speed.
+    def l2norm(self, sizea, sizeb, applier, repeat_tensors, in_type, per_tensor):
+        self.overflow_buf.zero_()
+        a = torch.cuda.FloatTensor(sizea).fill_(self.val)
+        b = torch.cuda.FloatTensor(sizeb).fill_(self.val)
+
+        in_list = []
+        for i in range(repeat_tensors):
+            in_list += [a.clone().to(in_type), b.clone().to(in_type)]
+
+        if per_tensor:
+            norm, norm_per_tensor = applier(multi_tensor_l2norm, self.overflow_buf, [in_list], True)
+            normab = torch.cat((a.norm().view(1), b.norm().view(1)))
+            norm_per_tensor = norm_per_tensor.view(-1, 2)
+        else:
+            norm, _ = applier(multi_tensor_l2norm, self.overflow_buf, [in_list], True)
+
+        reference = torch.cuda.FloatTensor((sizea + sizeb)*repeat_tensors).fill_(self.val).norm()
+
+        self.assertTrue(torch.allclose(norm, reference))
+        if per_tensor:
+          self.assertTrue(torch.allclose(norm_per_tensor, normab))
+        self.assertTrue(self.overflow_buf.item() == 0)
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_fuzz(self):
+        input_size_pairs = (
+            (7777*77, 555*555),
+            (777, 555),
+            (555, 2048*32+1),
+            (2048*32+1, 555),
+            (555, 2048*32),
+            (2048*32, 555),
+            (33333, 555),
+            (555, 33333))
+        appliers = (
+            MultiTensorApply(2048*32), 
+            MultiTensorApply(333),
+            MultiTensorApply(33333))
+        repeat_tensors = (
+            1,
+            55)
+
+        for sizea, sizeb in input_size_pairs:
+          for applier in appliers:
+            for repeat in repeat_tensors:
+              for in_type in (torch.float32, torch.float16):
+                for per_tensor in (False, True):
+                  self.l2norm(sizea, sizeb, applier, repeat, in_type, per_tensor)
+
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_scale.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_scale.py
new file mode 100644
index 0000000000..22da2490c6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multi_tensor_scale.py
@@ -0,0 +1,126 @@
+import unittest
+
+import functools as ft
+import itertools as it
+
+from apex import amp
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from utils import common_init, HALF, FLOAT,\
+    ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
+
+try:
+  import amp_C
+  from amp_C import multi_tensor_scale 
+  from apex.multi_tensor_apply import MultiTensorApply
+  disabled = False
+except ImportError as err:
+  print("amp_C fused kernels unavailable, disabling TestMultiTensorApply.  ImportError was ", err)
+  disabled = True
+
+
+class TestMultiTensorScale(unittest.TestCase):
+
+    def setUp(self):
+        common_init(self)
+        self.scale = 4.0
+        self.overflow_buf = torch.cuda.IntTensor(1).zero_()
+        self.ref = torch.cuda.FloatTensor([1.0])
+
+    def tearDown(self):
+        pass
+
+    # The tensor creation here is written for convenience, not speed.
+    def downscale(self, sizea, sizeb, applier, repeat_tensors, in_type, out_type, inplace=False):
+        self.overflow_buf.zero_()
+        a = torch.cuda.FloatTensor(sizea).fill_(self.scale)
+        b = torch.cuda.FloatTensor(sizeb).fill_(self.scale)
+
+        out_list = []
+        for i in range(repeat_tensors):
+            out_list += [a.clone().to(out_type), b.clone().to(out_type)]
+
+        if inplace:
+            in_list = out_list
+        else:
+            in_list = [out.clone().to(in_type) for out in out_list]
+
+        applier(multi_tensor_scale, self.overflow_buf, [in_list, out_list], 1./self.scale)
+
+        self.assertTrue(all([torch.allclose(out, self.ref.to(out_type)) for out in out_list]))
+        self.assertTrue(self.overflow_buf.item() == 0)
+ 
+    def find_inf(self, sizea, sizeb, applier, repeat_tensors, in_type, out_type, t, ind, val, inplace=False):
+        self.overflow_buf.zero_()
+        a = torch.cuda.FloatTensor(sizea).fill_(self.scale)
+        b = torch.cuda.FloatTensor(sizeb).fill_(self.scale)
+
+        out_list = []
+        for i in range(repeat_tensors):
+            out_list += [a.clone().to(out_type), b.clone().to(out_type)]
+
+        if inplace:
+            in_list = out_list
+        else:
+            in_list = [out.clone().to(in_type) for out in out_list]
+
+        applier(multi_tensor_scale, self.overflow_buf, [in_list, out_list], 1./self.scale)
+
+        self.overflow_buf.zero_()
+        in_list[t][ind] = val
+        applier(multi_tensor_scale, self.overflow_buf, [in_list, out_list], 1./self.scale)
+        self.assertTrue(self.overflow_buf.item())
+
+    # Currently, the fused kernel gives a hard error if you attempt to downscale
+    # into fp16 output, which imo is the desired behavior.  Maybe someday we
+    # will learn otherwise.
+    # @unittest.skipIf(disabled, "amp_C is unavailable")
+    # def test_fp16_to_fp16(self):
+    #     self.downscale(self.fp16, self.fp16, self.fp16_ref)
+    # 
+    # @unittest.skipIf(disabled, "amp_C is unavailable")
+    # def test_fp32_to_fp16(self):
+    #     self.downscale(self.fp32, self.fp16, self.fp16_ref)
+
+    @unittest.skipIf(disabled, "amp_C is unavailable")
+    def test_fuzz(self):
+        input_size_pairs = (
+            (7777*77, 555*555),
+            (777, 555),
+            (555, 2048*32+1),
+            (2048*32+1, 555),
+            (555, 2048*32),
+            (2048*32, 555),
+            (33333, 555),
+            (555, 33333))
+        appliers = (
+            MultiTensorApply(2048*32), 
+            MultiTensorApply(333),
+            MultiTensorApply(33333))
+        repeat_tensors = (
+            1,
+            55)
+
+        for sizea, sizeb in input_size_pairs:
+          for applier in appliers:
+            for repeat in repeat_tensors:
+              for in_type in (torch.float32, torch.float16):
+                for out_type in (torch.float32, torch.float16):
+                  for inplace in (True, False):
+                    if inplace is True and (out_type is not in_type):
+                      continue
+                    else:
+                      self.downscale(sizea, sizeb, applier, repeat, in_type, out_type, inplace=inplace)
+                      self.find_inf(sizea, sizeb, applier, repeat, in_type, out_type,
+                                    0, 0, float('nan'), inplace=inplace)
+                      self.find_inf(sizea, sizeb, applier, repeat, in_type, out_type,
+                                    2*repeat-1, sizeb-1, float('inf'), inplace=inplace)
+                      self.find_inf(sizea, sizeb, applier, repeat, in_type, out_type,
+                                   2*(repeat//2), sizea//2, float('inf'), inplace=inplace)
+
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multiple_models_optimizers_losses.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multiple_models_optimizers_losses.py
new file mode 100644
index 0000000000..068c845375
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_multiple_models_optimizers_losses.py
@@ -0,0 +1,762 @@
+import unittest
+
+import functools as ft
+import itertools as it
+
+from apex import amp
+from apex.amp import _amp_state
+import torch
+from torch import nn
+import torch.nn.functional as F
+from torch.nn import Parameter
+
+from utils import common_init, HALF, FLOAT,\
+    ALWAYS_HALF, ALWAYS_FLOAT, MATCH_INPUT
+
+class MyModel(torch.nn.Module):
+    def __init__(self, unique):
+        super(MyModel, self).__init__()
+        self.weight0 = Parameter(unique +
+            torch.arange(2, device='cuda', dtype=torch.float32))
+        self.weight1 = Parameter(1. + unique + torch.arange(2, device='cuda', dtype=torch.float16))
+
+    @staticmethod
+    def ops(input, weight0, weight1):
+        return ((input*(weight0.float()))*(weight1.float())).sum()
+
+    def forward(self, input):
+        return self.ops(input, self.weight0, self.weight1)
+
+# Abandon all hope, ye who enter here.
+
+# This is hands down the ugliest code I have ever written, but it succeeds in testing
+# multiple models/optimizers/losses fairly thoroughly.  Many of the different test cases
+# require slightly divergent code in a way that seems near-impossible to genericize into a simple
+# cross product or nested loops.
+
+class TestMultipleModelsOptimizersLosses(unittest.TestCase):
+    def setUp(self):
+        self.x = torch.ones((2), device='cuda', dtype=torch.float32)
+        common_init(self)
+
+    def tearDown(self):
+        pass
+
+    def test_2models2losses1optimizer(self):
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+
+        optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                     {'params' : model1.parameters(), 'lr' : 0.5}],
+                                    momentum=0.125)
+
+        reference_grads = []
+        for i in range(2):
+            optimizer.zero_grad()
+            loss0 = model0(self.x)
+            loss1 = model1(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads.append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()])
+
+            optimizer.step()
+
+        final_params = [param.data.clone() for param in model0.parameters()] + \
+                       [param.data.clone() for param in model1.parameters()]
+
+        for opt_level in ("O0", "O1", "O2", "O3"):
+          for how_to_zero in ("none", "model", "optimizer"):
+            for use_multiple_loss_scalers in (True, False):
+              if opt_level == "O1" or opt_level == "O2":
+                  inject_inf_iters = (-1, 0, 1)
+              else:
+                  inject_inf_iters = (-1,)
+
+              for inject_inf in inject_inf_iters:
+                if inject_inf >= 0:
+                   inject_inf_locs = ("fp16", "fp32")
+                   which_backwards = (0, 1)
+                else:
+                   inject_inf_locs = ("fdsa",)
+                   which_backwards = (None,)
+
+                for inject_inf_loc in inject_inf_locs:
+                  for which_backward in which_backwards:
+                      if use_multiple_loss_scalers:
+                          num_losses = 2
+                          loss_ids = [0, 1]
+                      else:
+                          num_losses = 1
+                          loss_ids = [0, 0]
+
+                      if inject_inf >= 0:
+                          iters = 3
+                      else:
+                          iters = 2
+
+                      model0 = MyModel(1)
+                      model1 = MyModel(2)
+
+                      models = [model0, model1]
+
+                      optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                                   {'params' : model1.parameters(), 'lr' : 0.5}],
+                                                  momentum=0.125)
+
+                      _amp_state.allow_incoming_model_not_fp32 = True
+                      [model0, model1], optimizer = amp.initialize(
+                          [model0, model1],
+                          optimizer,
+                          opt_level=opt_level,
+                          verbosity=0,
+                          cast_model_type=False,
+                          num_losses=num_losses)
+                      _amp_state.allow_incoming_model_not_fp32 = False
+
+                      _amp_state.loss_scalers[0]._loss_scale = 4.0
+                      if use_multiple_loss_scalers:
+                          _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                      unskipped = 0
+                      for i in range(iters):
+                          if how_to_zero == "none":
+                              for model in models:
+                                  for param in model.parameters():
+                                      param.grad = None
+                          elif how_to_zero == "model":
+                              for model in models:
+                                  model.zero_grad()
+                          else:
+                              optimizer.zero_grad()
+
+                          loss0 = model0(self.x)
+                          loss1 = model1(self.x)
+
+                          with amp.scale_loss(loss0, optimizer, loss_id=loss_ids[0]) as scaled_loss:
+                              scaled_loss.backward()
+                              if i == inject_inf and which_backward == 0:
+                                  if inject_inf_loc == "fp32":
+                                      model0.weight0.grad[0] = float('inf')
+                                  elif inject_inf_loc == "fp16":
+                                      model0.weight1.grad[0] = float('inf')
+                          with amp.scale_loss(loss1, optimizer, loss_id=loss_ids[1]) as scaled_loss:
+                              scaled_loss.backward()
+                              if i == inject_inf and which_backward == 1:
+                                  if inject_inf_loc == "fp32":
+                                      model1.weight0.grad[0] = float('inf')
+                                  elif inject_inf_loc == "fp16":
+                                      model1.weight1.grad[0] = float('inf')
+
+                          if i != inject_inf:
+                              for param, reference_grad in zip(amp.master_params(optimizer),
+                                                               reference_grads[unskipped]):
+                                  self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()))
+                              unskipped += 1
+                          optimizer.step()
+
+                      model_params = [p for p in model0.parameters()] + [p for p in model1.parameters()]
+                      for model, master, reference in zip(
+                              model_params,
+                              amp.master_params(optimizer),
+                              final_params):
+                          self.assertTrue(torch.allclose(model, reference))
+                          self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                      if opt_level == "O1":
+                          _amp_state.handle._deactivate()
+
+    def test_3models2losses1optimizer(self):
+
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+        model2 = MyModel(3)
+
+        optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                     {'params' : model1.parameters(), 'lr' : 0.5},
+                                     {'params' : model2.parameters(), 'lr' : 0.125}],
+                                     momentum=0.125)
+
+        reference_grads = []
+        for i in range(2):
+            optimizer.zero_grad()
+            loss0 = model0(self.x) + model2(self.x)
+            loss1 = model1(self.x) + model2(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads.append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()] +
+                                   [param.grad.data.clone() for param in model2.parameters()])
+
+            optimizer.step()
+
+
+        final_params = [param.data.clone() for param in model0.parameters()] + \
+                       [param.data.clone() for param in model1.parameters()] + \
+                       [param.data.clone() for param in model2.parameters()]
+
+        for opt_level in ("O0", "O1", "O2", "O3"):
+          for how_to_zero in ("none", "model", "optimizer"):
+            for use_multiple_loss_scalers in (True, False):
+              if opt_level == "O1" or opt_level == "O2":
+                  inject_inf_iters = (-1, 0, 1)
+              else:
+                  inject_inf_iters = (-1,)
+
+              for inject_inf in inject_inf_iters:
+                if inject_inf >= 0:
+                   inject_inf_locs = ("fp16", "fp32")
+                   which_backwards = (0, 1)
+                else:
+                   inject_inf_locs = ("fdsa",)
+                   which_backwards = (None,)
+
+                for inject_inf_loc in inject_inf_locs:
+                  for which_backward in which_backwards:
+                    if use_multiple_loss_scalers:
+                        num_losses = 2
+                        loss_ids = [0, 1]
+                    else:
+                        num_losses = 1
+                        loss_ids = [0, 0]
+
+                    if inject_inf >= 0:
+                        iters = 3
+                        if which_backward == 0:
+                            which_models = (0, 2)
+                        elif which_backward == 1:
+                            which_models = (1, 2)
+                    else:
+                        iters = 2
+                        which_models = (None,)
+
+                    for which_model in which_models:
+                        model0 = MyModel(1)
+                        model1 = MyModel(2)
+                        model2 = MyModel(3)
+
+                        models = [model0, model1, model2]
+
+                        optimizer = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                                     {'params' : model1.parameters(), 'lr' : 0.5},
+                                                     {'params' : model2.parameters(), 'lr' : 0.125}],
+                                                     momentum=0.125)
+
+                        _amp_state.allow_incoming_model_not_fp32 = True
+                        [model0, model1, model2], optimizer = amp.initialize(
+                            [model0, model1, model2],
+                            optimizer,
+                            opt_level=opt_level,
+                            verbosity=0,
+                            cast_model_type=False,
+                            num_losses=num_losses)
+                        _amp_state.allow_incoming_model_not_fp32 = False
+
+                        _amp_state.loss_scalers[0]._loss_scale = 4.0
+                        if use_multiple_loss_scalers:
+                            _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                        unskipped = 0
+                        for i in range(iters):
+                            if how_to_zero == "none":
+                                for model in models:
+                                    for param in model.parameters():
+                                        param.grad = None
+                            elif how_to_zero == "model":
+                                for model in models:
+                                    model.zero_grad()
+                            else:
+                                optimizer.zero_grad()
+
+                            # print("opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} which_model {} use_multiple_loss_scalers {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, which_model, use_multiple_loss_scalers))
+
+                            loss0 = model0(self.x) + model2(self.x)
+                            loss1 = model1(self.x) + model2(self.x)
+
+                            with amp.scale_loss(loss0, optimizer, loss_id=loss_ids[0]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 0:
+                                    if which_model == 0:
+                                        inj_model = model0
+                                    elif which_model == 2:
+                                        inj_model = model2
+                                    else:
+                                        raise RuntimeError(which_model + " invalid for loss 0")
+                                    if inject_inf_loc == "fp32":
+                                        inj_model.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        inj_model.weight1.grad[0] = float('inf')
+                            with amp.scale_loss(loss1, optimizer, loss_id=loss_ids[1]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 1:
+                                    if which_model == 1:
+                                        inj_model = model1
+                                    elif which_model == 2:
+                                        inj_model = model2
+                                    else:
+                                        raise RuntimeError(which_model + " invalid for loss 1 ")
+                                    if inject_inf_loc == "fp32":
+                                        inj_model.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        inj_model.weight1.grad[0] = float('inf')
+
+                            if i != inject_inf:
+                                for param, reference_grad in zip(amp.master_params(optimizer),
+                                                                 reference_grads[unskipped]):
+                                    self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()))
+                                unskipped += 1
+
+                            optimizer.step()
+
+                        model_params = [p for p in model0.parameters()] + \
+                                       [p for p in model1.parameters()] + \
+                                       [p for p in model2.parameters()]
+                        for model, master, reference in zip(
+                                model_params,
+                                amp.master_params(optimizer),
+                                final_params):
+                            self.assertTrue(torch.allclose(model, reference))
+                            self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                        if opt_level == "O1":
+                            _amp_state.handle._deactivate()
+
+    def test_2models2losses2optimizers(self):
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+
+        optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                      momentum=0.125)
+        optimizer1 = torch.optim.SGD([{'params' : model1.parameters(), 'lr' : 0.5}],
+                                      momentum=0.25)
+
+        # Don't do it like this:  reference_grads = [[]]*5
+        # because then it creates a list of 5 references to the same "[]" and appending
+        # to any of them effectively makes you append to all of them, which multiplies
+        # the resulting size of reference_grads by 5x and needless to say makes the test fail.
+        reference_grads = [[], [], [], [], []]
+        final_params = [None, None, None, None, None]
+        for i in range(2):
+            optimizer0.zero_grad()
+            optimizer1.zero_grad()
+            loss0 = model0(self.x)
+            loss1 = model1(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads[0].append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()])
+
+            optimizer0.step()
+            optimizer1.step()
+
+        final_params[0] = [param.data.clone() for param in model0.parameters()] + \
+                          [param.data.clone() for param in model1.parameters()]
+
+        def what_got_skipped(which_iter, which_backward):
+            if which_iter == 0 and which_backward == 0:
+                return 1
+            if which_iter == 0 and which_backward == 1:
+                return 2
+            if which_iter == 1 and which_backward == 0:
+                return 3
+            if which_iter == 1 and which_backward == 1:
+                return 4
+            return 0
+
+        for which_iter in (0,1):
+            for which_backward in (0,1):
+                model0 = MyModel(1)
+                model1 = MyModel(2)
+
+                optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                              momentum=0.125)
+                optimizer1 = torch.optim.SGD([{'params' : model1.parameters(), 'lr' : 0.5}],
+                                              momentum=0.25)
+
+                for i in range(3):
+                    optimizer0.zero_grad()
+                    optimizer1.zero_grad()
+                    loss0 = model0(self.x)
+                    loss1 = model1(self.x)
+                    loss0.backward()
+                    loss1.backward()
+
+                    if i != which_iter:
+                        reference_grads[what_got_skipped(which_iter, which_backward)].append(
+                            [param.grad.data.clone() for param in model0.parameters()] +
+                            [param.grad.data.clone() for param in model1.parameters()])
+
+                    if i == which_iter:
+                        if which_backward == 0:
+                            optimizer1.step()
+                        else:
+                            optimizer0.step()
+                    else:
+                        optimizer0.step()
+                        optimizer1.step()
+
+                final_params[what_got_skipped(which_iter, which_backward)] = \
+                    [param.data.clone() for param in model0.parameters()] + \
+                    [param.data.clone() for param in model1.parameters()]
+
+        for opt_level in ("O0", "O1", "O2", "O3"):
+          for how_to_zero in ("none", "model", "optimizer"):
+            for use_multiple_loss_scalers in (True, False):
+              if opt_level == "O1" or opt_level == "O2":
+                  inject_inf_iters = (-1, 0, 1)
+              else:
+                  inject_inf_iters = (-1,)
+
+              for inject_inf in inject_inf_iters:
+                if inject_inf >= 0:
+                   inject_inf_locs = ("fp16", "fp32")
+                   which_backwards = (0, 1)
+                else:
+                   inject_inf_locs = ("fdsa",)
+                   which_backwards = (None,)
+
+                for inject_inf_loc in inject_inf_locs:
+                  for which_backward in which_backwards:
+                      if use_multiple_loss_scalers:
+                          num_losses = 2
+                          loss_ids = [0, 1]
+                      else:
+                          num_losses = 1
+                          loss_ids = [0, 0]
+
+                      if inject_inf >= 0:
+                          iters = 3
+                      else:
+                          iters = 2
+
+                      model0 = MyModel(1)
+                      model1 = MyModel(2)
+
+                      models = [model0, model1]
+
+                      optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25}],
+                                                    momentum=0.125)
+                      optimizer1 = torch.optim.SGD([{'params' : model1.parameters(), 'lr' : 0.5}],
+                                                    momentum=0.25)
+
+                      _amp_state.allow_incoming_model_not_fp32 = True
+                      [model0, model1], [optimizer0, optimizer1] = amp.initialize(
+                          [model0, model1],
+                          [optimizer0, optimizer1],
+                          opt_level=opt_level,
+                          verbosity=0,
+                          cast_model_type=False,
+                          num_losses=num_losses)
+                      _amp_state.allow_incoming_model_not_fp32 = False
+
+                      _amp_state.loss_scalers[0]._loss_scale = 4.0
+                      if use_multiple_loss_scalers:
+                          _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                      unskipped = 0
+                      for i in range(iters):
+                          if how_to_zero == "none":
+                              for model in models:
+                                  for param in model.parameters():
+                                      param.grad = None
+                          elif how_to_zero == "model":
+                              for model in models:
+                                  model.zero_grad()
+                          else:
+                              optimizer0.zero_grad()
+                              optimizer1.zero_grad()
+
+                          loss0 = model0(self.x)
+                          loss1 = model1(self.x)
+
+                          with amp.scale_loss(loss0, optimizer0, loss_id=loss_ids[0]) as scaled_loss:
+                              scaled_loss.backward()
+                              if i == inject_inf and which_backward == 0:
+                                  if inject_inf_loc == "fp32":
+                                      model0.weight0.grad[0] = float('inf')
+                                  elif inject_inf_loc == "fp16":
+                                      model0.weight1.grad[0] = float('inf')
+                          with amp.scale_loss(loss1, optimizer1, loss_id=loss_ids[1]) as scaled_loss:
+                              scaled_loss.backward()
+                              if i == inject_inf and which_backward == 1:
+                                  if inject_inf_loc == "fp32":
+                                      model1.weight0.grad[0] = float('inf')
+                                  elif inject_inf_loc == "fp16":
+                                      model1.weight1.grad[0] = float('inf')
+
+                          # print("opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} use_multiple_loss_scalers {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, use_multiple_loss_scalers))
+
+                          if i != inject_inf:
+                              master_params = list(amp.master_params(optimizer0)) + \
+                                              list(amp.master_params(optimizer1))
+                              for param, reference_grad in zip(master_params,
+                                      reference_grads[what_got_skipped(inject_inf, which_backward)][unskipped]):
+                                  self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()))
+                              unskipped += 1
+
+                          optimizer0.step()
+                          optimizer1.step()
+
+                      model_params = [p for p in model0.parameters()] + [p for p in model1.parameters()]
+                      master_params = [p for p in amp.master_params(optimizer0)] + \
+                                      [p for p in amp.master_params(optimizer1)]
+                      for model, master, reference in zip(
+                              model_params,
+                              master_params,
+                              final_params[what_got_skipped(inject_inf, which_backward)]):
+                          self.assertTrue(torch.allclose(model, reference))
+                          self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                      if opt_level == "O1":
+                          _amp_state.handle._deactivate()
+
+    def test_3models2losses2optimizers(self):
+        model0 = MyModel(1)
+        model1 = MyModel(2)
+        model2 = MyModel(3)
+
+        optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                      {'params' : model1.parameters(), 'lr' : 1.0}],
+                                     momentum=0.5)
+        optimizer1 = torch.optim.SGD([{'params' : model2.parameters(), 'lr' : 0.5}],
+                                     momentum=0.25)
+
+        # Again, can't do this:  reference_grads = [[]]*9
+        reference_grads = [[], [], [], [], [], [], [], [], []]
+        final_params = [None, None, None, None, None, None, None, None, None]
+        for i in range(2):
+            optimizer0.zero_grad()
+            optimizer1.zero_grad()
+            loss0 = model0(self.x) + model1(self.x)
+            loss1 = model2(self.x) + model1(self.x)
+            loss0.backward()
+            loss1.backward()
+
+            reference_grads[0].append([param.grad.data.clone() for param in model0.parameters()] +
+                                   [param.grad.data.clone() for param in model1.parameters()])
+
+            optimizer0.step()
+            optimizer1.step()
+
+        final_params[0] = \
+            [param.data.clone() for param in model0.parameters()] + \
+            [param.data.clone() for param in model1.parameters()] + \
+            [param.data.clone() for param in model2.parameters()]
+
+        def what_got_skipped(which_iter, which_backward, which_model):
+            if which_iter == 0:
+                if which_backward == 0:
+                    if which_model == 0:
+                        return 1
+                    if which_model == 1:
+                        return 2
+                if which_backward == 1:
+                    if which_model == 2:
+                        return 3
+                    if which_model == 1:
+                        return 4
+            if which_iter == 1:
+                if which_backward == 0:
+                    if which_model == 0:
+                        return 5
+                    if which_model == 1:
+                        return 6
+                if which_backward == 1:
+                    if which_model == 2:
+                        return 7
+                    if which_model == 1:
+                        return 8
+            return 0
+
+        for which_iter in (0,1):
+            for which_backward in (0,1):
+                if which_backward == 0:
+                    which_models = (0,1)
+                if which_backward == 1:
+                    which_models = (2,1)
+                for which_model in which_models:
+
+                    model0 = MyModel(1)
+                    model1 = MyModel(2)
+                    model2 = MyModel(3)
+
+                    optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                                  {'params' : model1.parameters(), 'lr' : 1.0}],
+                                                 momentum=0.5)
+                    optimizer1 = torch.optim.SGD([{'params' : model2.parameters(), 'lr' : 0.5}],
+                                                 momentum=0.25)
+
+                    for i in range(3):
+                        optimizer0.zero_grad()
+                        optimizer1.zero_grad()
+                        loss0 = model0(self.x) + model1(self.x)
+                        loss1 = model2(self.x) + model1(self.x)
+                        loss0.backward()
+                        loss1.backward()
+
+                        if i != which_iter:
+                            reference_grads[what_got_skipped(which_iter,
+                                    which_backward, which_model)].append(
+                                [param.grad.data.clone() for param in model0.parameters()] +
+                                [param.grad.data.clone() for param in model1.parameters()])
+
+                        if i == which_iter:
+                            if which_backward == 0:
+                                # if which_model == 0:
+                                    optimizer1.step()
+                                # if which_model == 1:
+                                #     optimizer1.step()
+                            if which_backward == 1:
+                                # if which_model == 2:
+                                #     optimizer0.step()
+                                # if which_model == 1:
+                                    continue
+                        else:
+                            optimizer0.step()
+                            optimizer1.step()
+
+                    final_params[what_got_skipped(which_iter, which_backward, which_model)] = \
+                        [param.data.clone() for param in model0.parameters()] + \
+                        [param.data.clone() for param in model1.parameters()] + \
+                        [param.data.clone() for param in model2.parameters()]
+
+        for opt_level in ("O0", "O1", "O2", "O3"):
+          for how_to_zero in ("none", "model", "optimizer"):
+            for use_multiple_loss_scalers in (True, False):
+              if opt_level == "O1" or opt_level == "O2":
+                  inject_inf_iters = (-1, 0, 1)
+              else:
+                  inject_inf_iters = (-1,)
+
+              for inject_inf in inject_inf_iters:
+                if inject_inf >= 0:
+                   inject_inf_locs = ("fp16", "fp32")
+                   which_backwards = (0, 1)
+                else:
+                   inject_inf_locs = ("fdsa",)
+                   which_backwards = (None,)
+
+                for inject_inf_loc in inject_inf_locs:
+                  for which_backward in which_backwards:
+                    if use_multiple_loss_scalers:
+                        num_losses = 2
+                        loss_ids = [0, 1]
+                    else:
+                        num_losses = 1
+                        loss_ids = [0, 0]
+
+                    if inject_inf >= 0:
+                        iters = 3
+                        if which_backward == 0:
+                            which_models = (0, 1)
+                        elif which_backward == 1:
+                            which_models = (2, 1)
+                    else:
+                        iters = 2
+                        which_models = (None,)
+
+                    for which_model in which_models:
+                        model0 = MyModel(1)
+                        model1 = MyModel(2)
+                        model2 = MyModel(3)
+
+                        models = [model0, model1, model2]
+
+                        optimizer0 = torch.optim.SGD([{'params' : model0.parameters(), 'lr' : 0.25},
+                                                      {'params' : model1.parameters(), 'lr' : 1.0}],
+                                                     momentum=0.5)
+                        optimizer1 = torch.optim.SGD([{'params' : model2.parameters(), 'lr' : 0.5}],
+                                                     momentum=0.25)
+
+                        _amp_state.allow_incoming_model_not_fp32 = True
+                        [model0, model1, model2], [optimizer0, optimizer1] = amp.initialize(
+                            [model0, model1, model2],
+                            [optimizer0, optimizer1],
+                            opt_level=opt_level,
+                            verbosity=0,
+                            cast_model_type=False,
+                            num_losses=num_losses)
+                        _amp_state.allow_incoming_model_not_fp32 = False
+
+                        _amp_state.loss_scalers[0]._loss_scale = 4.0
+                        if use_multiple_loss_scalers:
+                            _amp_state.loss_scalers[1]._loss_scale = 16.0
+
+                        unskipped = 0
+                        for i in range(iters):
+                            if how_to_zero == "none":
+                                for model in models:
+                                    for param in model.parameters():
+                                        param.grad = None
+                            elif how_to_zero == "model":
+                                for model in models:
+                                    model.zero_grad()
+                            else:
+                                optimizer0.zero_grad()
+                                optimizer1.zero_grad()
+
+                            loss0 = model0(self.x) + model1(self.x)
+                            loss1 = model2(self.x) + model1(self.x)
+
+                            with amp.scale_loss(loss0, optimizer0, loss_id=loss_ids[0]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 0:
+                                    if which_model == 0:
+                                        inj_model = model0
+                                    elif which_model == 1:
+                                        inj_model = model1
+                                    else:
+                                        raise RuntimeError(which_model + " invalid for loss 0")
+                                    if inject_inf_loc == "fp32":
+                                        inj_model.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        inj_model.weight1.grad[0] = float('inf')
+                            with amp.scale_loss(loss1, [optimizer0, optimizer1], loss_id=loss_ids[1]) as scaled_loss:
+                                scaled_loss.backward()
+                                if i == inject_inf and which_backward == 1:
+                                    if which_model == 2:
+                                        inj_model = model2
+                                    elif which_model == 1:
+                                        inj_model = model1
+                                    else:
+                                        raise RuntimeError(which_model + " invalid for loss 1 ")
+                                    if inject_inf_loc == "fp32":
+                                        inj_model.weight0.grad[0] = float('inf')
+                                    elif inject_inf_loc == "fp16":
+                                        inj_model.weight1.grad[0] = float('inf')
+
+                            if i != inject_inf:
+                                master_params = list(amp.master_params(optimizer0)) + \
+                                                list(amp.master_params(optimizer1))
+                                for param, reference_grad in zip(master_params,
+                                      reference_grads[what_got_skipped(inject_inf,
+                                          which_backward, which_model)][unskipped]):
+                                    self.assertTrue(torch.allclose(param.grad.float(), reference_grad.float()))
+                                unskipped += 1
+
+                            optimizer0.step()
+                            optimizer1.step()
+
+                        model_params = [p for p in model0.parameters()] + \
+                                       [p for p in model1.parameters()] + \
+                                       [p for p in model2.parameters()]
+                        master_params = [p for p in amp.master_params(optimizer0)] + \
+                                        [p for p in amp.master_params(optimizer1)]
+
+                        # print("opt_level {} i {} inject_inf {} which_backward {} inject_inf_loc {} use_multiple_loss_scalers {} which_model {}".format(opt_level, i, inject_inf, which_backward, inject_inf_loc, use_multiple_loss_scalers, which_model))
+
+                        for model, master, reference in zip(
+                                model_params,
+                                master_params,
+                                final_params[what_got_skipped(inject_inf, which_backward, which_model)]):
+                            self.assertTrue(torch.allclose(model, reference))
+                            self.assertTrue(torch.allclose(model, master.to(model.dtype)))
+
+                        if opt_level == "O1":
+                            _amp_state.handle._deactivate()
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_promotion.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_promotion.py
new file mode 100644
index 0000000000..f5ef30c127
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_promotion.py
@@ -0,0 +1,75 @@
+import unittest
+
+import itertools as it
+
+from apex import amp
+import torch
+from torch import nn
+import torch.nn.functional as F
+
+from utils import common_init, HALF, FLOAT, DTYPES
+
+class TestPromotion(unittest.TestCase):
+    def setUp(self):
+        self.handle = amp.init(enabled=True)
+        common_init(self)
+
+    def tearDown(self):
+        self.handle._deactivate()
+
+    def run_binary_promote_test(self, fns, input_shape, x_inplace=False):
+        type_pairs = it.product(DTYPES, DTYPES)
+        for fn, (xtype, ytype) in it.product(fns, type_pairs):
+            x = torch.randn(input_shape, dtype=xtype).requires_grad_()
+            x_leaf = x
+            if x_inplace:
+                # We need a non-leaf to call in place on
+                x = x.clone()
+            y = torch.randn(input_shape, dtype=ytype)
+            out = fn(x, y)
+            if x_inplace:
+                # In place: always match xtype
+                self.assertEqual(out.type(), x.type())
+            else:
+                # Out of place: match widest type
+                if xtype == torch.float or ytype == torch.float:
+                    self.assertEqual(out.type(), FLOAT)
+                else:
+                    self.assertEqual(out.type(), HALF)
+            out.float().sum().backward()
+            self.assertEqual(x_leaf.grad.dtype, xtype)
+
+    def test_atan2_matches_widest(self):
+        fns = [lambda x, y : torch.atan2(x, y),
+               lambda x, y : x.atan2(y)]
+        self.run_binary_promote_test(fns, (self.b,))
+
+    def test_mul_matches_widest(self):
+        fns = [lambda x, y : torch.mul(x, y),
+               lambda x, y: x.mul(y)]
+        self.run_binary_promote_test(fns, (self.b,))
+
+    def test_cat_matches_widest(self):
+        shape = self.b
+        ys = [torch.randn(shape, dtype=torch.half) for _ in range(5)]
+        x_float = torch.randn(shape)
+        out = torch.cat(ys + [x_float])
+        self.assertEqual(out.type(), FLOAT)
+        x_half = torch.randn(shape, dtype=torch.half)
+        out = torch.cat(ys + [x_half])
+        self.assertEqual(out.type(), HALF)
+
+    def test_inplace_exp_is_error_for_half(self):
+        xs = torch.randn(self.b)
+        xs.exp_()
+        self.assertEqual(xs.type(), FLOAT)
+        xs = torch.randn(self.b, dtype=torch.half)
+        with self.assertRaises(NotImplementedError):
+            xs.exp_()
+
+    def test_inplace_add_matches_self(self):
+        fn = lambda x, y: x.add_(y)
+        self.run_binary_promote_test([fn], (self.b,), x_inplace=True)
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_rnn.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_rnn.py
new file mode 100644
index 0000000000..c49a5f0034
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/test_rnn.py
@@ -0,0 +1,116 @@
+import unittest
+
+from apex import amp
+import random
+import torch
+from torch import nn
+
+from utils import common_init, HALF
+
+class TestRnnCells(unittest.TestCase):
+    def setUp(self):
+        self.handle = amp.init(enabled=True)
+        common_init(self)
+
+    def tearDown(self):
+        self.handle._deactivate()
+
+    def run_cell_test(self, cell, state_tuple=False):
+        shape = (self.b, self.h)
+        for typ in [torch.float, torch.half]:
+            xs = [torch.randn(shape, dtype=typ).requires_grad_()
+                  for _ in range(self.t)]
+            hidden_fn = lambda: torch.zeros(shape, dtype=typ)
+            if state_tuple:
+                hidden = (hidden_fn(), hidden_fn())
+            else:
+                hidden = hidden_fn()
+            outputs = []
+            for i in range(self.t):
+                hidden = cell(xs[i], hidden)
+                if state_tuple:
+                    output = hidden[0]
+                else:
+                    output = hidden
+                outputs.append(output)
+            for y in outputs:
+                self.assertEqual(y.type(), HALF)
+            outputs[-1].float().sum().backward()
+            for i, x in enumerate(xs):
+                self.assertEqual(x.grad.dtype, x.dtype)
+
+    def test_rnn_cell_is_half(self):
+        cell = nn.RNNCell(self.h, self.h)
+        self.run_cell_test(cell)
+
+    def test_gru_cell_is_half(self):
+        cell = nn.GRUCell(self.h, self.h)
+        self.run_cell_test(cell)
+
+    def test_lstm_cell_is_half(self):
+        cell = nn.LSTMCell(self.h, self.h)
+        self.run_cell_test(cell, state_tuple=True)
+
+class TestRnns(unittest.TestCase):
+    def setUp(self):
+        self.handle = amp.init(enabled=True)
+        common_init(self)
+
+    def tearDown(self):
+        self.handle._deactivate()
+
+    def run_rnn_test(self, rnn, layers, bidir, state_tuple=False):
+        for typ in [torch.float, torch.half]:
+            x = torch.randn((self.t, self.b, self.h), dtype=typ).requires_grad_()
+            hidden_fn = lambda: torch.zeros((layers + (layers * bidir),
+                                             self.b, self.h), dtype=typ)
+            if state_tuple:
+                hidden = (hidden_fn(), hidden_fn())
+            else:
+                hidden = hidden_fn()
+            output, _ = rnn(x, hidden)
+            self.assertEqual(output.type(), HALF)
+            output[-1, :, :].float().sum().backward()
+            self.assertEqual(x.grad.dtype, x.dtype)
+
+    def test_rnn_is_half(self):
+        configs = [(1, False), (2, False), (2, True)]
+        for layers, bidir in configs:
+            rnn = nn.RNN(input_size=self.h, hidden_size=self.h, num_layers=layers,
+                         nonlinearity='relu', bidirectional=bidir)
+            self.run_rnn_test(rnn, layers, bidir)
+
+    def test_gru_is_half(self):
+        configs = [(1, False), (2, False), (2, True)]
+        for layers, bidir in configs:
+            rnn = nn.GRU(input_size=self.h, hidden_size=self.h, num_layers=layers,
+                         bidirectional=bidir)
+            self.run_rnn_test(rnn, layers, bidir)
+
+    def test_lstm_is_half(self):
+        configs = [(1, False), (2, False), (2, True)]
+        for layers, bidir in configs:
+            rnn = nn.LSTM(input_size=self.h, hidden_size=self.h, num_layers=layers,
+                         bidirectional=bidir)
+            self.run_rnn_test(rnn, layers, bidir, state_tuple=True)
+
+    def test_rnn_packed_sequence(self):
+        num_layers = 2
+        rnn = nn.RNN(input_size=self.h, hidden_size=self.h, num_layers=num_layers)
+        for typ in [torch.float, torch.half]:
+            x = torch.randn((self.t, self.b, self.h), dtype=typ).requires_grad_()
+            lens = sorted([random.randint(self.t // 2, self.t) for _ in range(self.b)],
+                          reverse=True)
+            # `pack_padded_sequence` breaks if default tensor type is non-CPU
+            torch.set_default_tensor_type(torch.FloatTensor)
+            lens = torch.tensor(lens, dtype=torch.int64, device=torch.device('cpu'))
+            packed_seq = nn.utils.rnn.pack_padded_sequence(x, lens)
+            torch.set_default_tensor_type(torch.cuda.FloatTensor)
+            hidden = torch.zeros((num_layers, self.b, self.h), dtype=typ)
+            output, _ = rnn(packed_seq, hidden)
+            self.assertEqual(output.data.type(), HALF)
+            output.data.float().sum().backward()
+            self.assertEqual(x.grad.dtype, x.dtype)
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/utils.py
new file mode 100644
index 0000000000..7aa20c3691
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_amp/utils.py
@@ -0,0 +1,21 @@
+import torch
+
+HALF = 'torch.cuda.HalfTensor'
+FLOAT = 'torch.cuda.FloatTensor'
+
+DTYPES = [torch.half, torch.float]
+
+ALWAYS_HALF = {torch.float: HALF,
+               torch.half: HALF}
+ALWAYS_FLOAT = {torch.float: FLOAT,
+                torch.half: FLOAT}
+MATCH_INPUT = {torch.float: FLOAT,
+               torch.half: HALF}
+
+def common_init(test_case):
+    test_case.h = 64
+    test_case.b = 16
+    test_case.c = 16
+    test_case.k = 3
+    test_case.t = 10
+    torch.set_default_tensor_type(torch.cuda.FloatTensor)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fp16util/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fp16util/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fp16util/test_fp16util.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fp16util/test_fp16util.py
new file mode 100644
index 0000000000..eecddbc01d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fp16util/test_fp16util.py
@@ -0,0 +1,75 @@
+import unittest
+
+import torch
+import torch.nn as nn
+
+from apex.fp16_utils import FP16Model
+
+
+class DummyBlock(nn.Module):
+    def __init__(self):
+        super(DummyBlock, self).__init__()
+
+        self.conv = nn.Conv2d(10, 10, 2)
+        self.bn = nn.BatchNorm2d(10, affine=True)
+
+    def forward(self, x):
+        return self.conv(self.bn(x))
+
+
+class DummyNet(nn.Module):
+    def __init__(self):
+        super(DummyNet, self).__init__()
+
+        self.conv1 = nn.Conv2d(3, 10, 2)
+        self.bn1 = nn.BatchNorm2d(10, affine=False)
+        self.db1 = DummyBlock()
+        self.db2 = DummyBlock()
+
+    def forward(self, x):
+        out = x
+        out = self.conv1(out)
+        out = self.bn1(out)
+        out = self.db1(out)
+        out = self.db2(out)
+        return out
+
+
+class DummyNetWrapper(nn.Module):
+    def __init__(self):
+        super(DummyNetWrapper, self).__init__()
+
+        self.bn = nn.BatchNorm2d(3, affine=True)
+        self.dn = DummyNet()
+
+    def forward(self, x):
+        return self.dn(self.bn(x))
+
+
+class TestFP16Model(unittest.TestCase):
+    def setUp(self):
+        self.N = 64
+        self.C_in = 3
+        self.H_in = 16
+        self.W_in = 32
+        self.in_tensor = torch.randn((self.N, self.C_in, self.H_in, self.W_in)).cuda()
+        self.orig_model = DummyNetWrapper().cuda()
+        self.fp16_model = FP16Model(self.orig_model)
+
+    def test_params_and_buffers(self):
+        exempted_modules = [
+            self.fp16_model.network.bn,
+            self.fp16_model.network.dn.db1.bn,
+            self.fp16_model.network.dn.db2.bn,
+        ]
+        for m in self.fp16_model.modules():
+            expected_dtype = torch.float if (m in exempted_modules) else torch.half
+            for p in m.parameters(recurse=False):
+                assert p.dtype == expected_dtype
+            for b in m.buffers(recurse=False):
+                assert b.dtype in (expected_dtype, torch.int64)
+
+    def test_output_is_half(self):
+        out_tensor = self.fp16_model(self.in_tensor)
+        assert out_tensor.dtype == torch.half
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py
new file mode 100644
index 0000000000..67fb8b2bd8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_fused_layer_norm/test_fused_layer_norm.py
@@ -0,0 +1,143 @@
+import itertools
+import unittest
+
+import torch
+
+import apex
+
+
+class TestFusedLayerNorm(unittest.TestCase):
+    dtype = torch.float
+    elementwise_affine = False
+    normalized_shape = [32, 16]
+    rtol, atol = None, None
+    fwd_thresholds = dict(rtol=None, atol=None)
+    bwd_thresholds = dict(rtol=None, atol=None)
+
+    def setUp(self):
+        # bias and weight are set to 0 and 1 respectively, so no need to copy parameters from cpu module to the gpu one
+        self.module_cpu_ = apex.normalization.FusedLayerNorm(
+            normalized_shape=self.normalized_shape, elementwise_affine=self.elementwise_affine).cpu()
+        self.module_cuda_ = apex.normalization.FusedLayerNorm(
+            normalized_shape=self.normalized_shape, elementwise_affine=self.elementwise_affine).to(device="cuda", dtype=self.dtype)
+
+    def _check_same_output(self, batch_size, contiguous):
+        torch.cuda.manual_seed(42)
+        if contiguous:
+            input_shape = [batch_size] + self.normalized_shape
+            input_ = torch.randn(input_shape, device="cpu").requires_grad_(True)
+            input_cuda_ = input_.to(device="cuda", dtype=self.dtype).detach().requires_grad_(True)
+            self.assertTrue(input_.is_contiguous())
+            self.assertTrue(input_cuda_.is_contiguous())
+        else:
+            input_shape = [batch_size] + self.normalized_shape
+            input_shape = [batch_size * 3] + [self.normalized_shape[0] * 5, self.normalized_shape[1] * 3]
+            input_src_ = torch.randn(input_shape, device="cpu")
+            input_ = input_src_[::3, ::5, ::3].detach().requires_grad_(True)
+            input_cuda_ = input_src_.to(device="cuda", dtype=self.dtype)[::3, ::5, ::3].detach().requires_grad_(True)
+            # make sure that tensors are NOT contiguous.
+            self.assertFalse(input_.is_contiguous())
+            self.assertFalse(input_cuda_.is_contiguous())
+        out_cpu_ = self.module_cpu_(input_)
+        gO = torch.rand_like(out_cpu_)
+        out_cpu_.backward(gO)
+        out_cuda_ = self.module_cuda_(input_cuda_)
+        gO = gO.to(device="cuda", dtype=self.dtype)
+        out_cuda_.backward(gO)
+        self.assertFalse(out_cpu_.is_cuda)
+        self.assertTrue(out_cuda_.is_cuda)
+        # TODO (mkozuki): `torch.testing.assert_allclose` is deprecated.
+        # Use `torch.testing.assert_close`.
+        # See https://github.com/pytorch/pytorch/issues/61844
+        torch.testing.assert_allclose(
+            out_cpu_.to(device="cuda", dtype=self.dtype), out_cuda_, **self.fwd_thresholds)
+        torch.testing.assert_allclose(
+            input_.grad.to(device="cuda", dtype=self.dtype), input_cuda_.grad, **self.bwd_thresholds)
+
+    def _test_same_output(self, batch_size):
+        for contiguous in (True, False):
+            with self.subTest(contiguous=contiguous):
+                self._check_same_output(batch_size, contiguous)
+
+    def test_layer_norm(self):
+        self._test_same_output(16)
+
+    def test_large_batch(self):
+        self._test_same_output(65536)
+
+
+class TestFusedLayerNormElemWise(TestFusedLayerNorm):
+    elementwise_affine = True
+
+
+class TestFusedLayerNormElemWiseHalf(TestFusedLayerNormElemWise):
+    dtype = torch.half
+
+    def test_large_batch(self):
+        self.skipTest("Skip to save time")
+
+
+class TestFusedLayerNormElemWiseBFloat16(TestFusedLayerNormElemWise):
+    dtype = torch.bfloat16
+    # NOTE (mkozuki): [BFloat16 Layer Norm flakiness]
+    # Use thresholds larger than those used in pytorch, see
+    # https://github.com/pytorch/pytorch/blob/72274e2a2fd55019ec860e1743dbdc5b0c5a5624/torch/testing/_asserts.py#L26
+    fwd_thresholds = dict(rtol=1.6e-2, atol=3e-4)
+    bwd_thresholds = dict(rtol=1.6e-2, atol=3e-3)
+
+    def test_large_batch(self):
+        self.skipTest("Skip to save time")
+
+
+def _prep_layers(normalized_shape, elementwise_affine, dtype):
+    native = torch.nn.LayerNorm(
+        normalized_shape=normalized_shape, elementwise_affine=elementwise_affine
+    ).to(device="cuda", dtype=dtype)
+    fused = apex.normalization.FusedLayerNorm(
+        normalized_shape=normalized_shape, elementwise_affine=elementwise_affine
+    ).cuda()
+    return native, fused
+
+
+def _prep_inputs(batch_size, normalized_shape, dtype):
+    shape = (batch_size, *normalized_shape)
+    fused = torch.randn(shape).cuda().requires_grad_(True)
+    with torch.no_grad():
+        native = fused.clone().to(dtype).requires_grad_(True)
+    return native, fused
+
+
+autocast_dtypes = (torch.half, torch.bfloat16) if torch.cuda.is_bf16_supported() else (torch.half,)
+
+
+class TestAutocastFusedLayerNorm(unittest.TestCase):
+    bf16_fwd_thresholds = dict(rtol=1.6e-2, atol=3e-4)
+    bf16_bwd_thresholds = dict(rtol=1.6e-2, atol=3e-3)
+
+    def setUp(self):
+        self.batch_size = 16
+        self.normalized_shape = [32, 16]
+
+    def _run_test(self, dtype, elementwise_affine):
+        native, fused = _prep_layers(self.normalized_shape, elementwise_affine, dtype)
+        native_x, fused_x = _prep_inputs(self.batch_size, self.normalized_shape, dtype)
+
+        expected = native(native_x)
+        with torch.cuda.amp.autocast(dtype=dtype):
+            actual = fused(fused_x)
+        tols = {'rtol': None, 'atol': None} if dtype == torch.half else TestAutocastFusedLayerNorm.bf16_fwd_thresholds
+        torch.testing.assert_allclose(actual, expected, **tols)
+
+        g_native = torch.rand_like(expected)
+        with torch.no_grad():
+            g_fused = g_native.clone()
+        expected.backward(g_native)
+        actual.backward(g_fused)
+
+        tols = {'rtol': None, 'atol': None} if dtype == torch.half else TestAutocastFusedLayerNorm.bf16_bwd_thresholds
+        torch.testing.assert_allclose(native_x.grad, fused_x.grad, **tols)
+
+    def test_autocast(self):
+        for (dtype, elementwise_affine) in itertools.product(autocast_dtypes, (True, False)):
+            with self.subTest(f"{dtype}-{elementwise_affine}"):
+                self._run_test(dtype, elementwise_affine)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_mlp/test_mlp.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_mlp/test_mlp.py
new file mode 100644
index 0000000000..9ccda566dd
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_mlp/test_mlp.py
@@ -0,0 +1,218 @@
+"""Tests for c++ MLP"""
+import unittest
+from time import time
+import numpy as np
+
+import torch
+from torch import nn
+
+from apex.mlp import MLP
+
+batch_size = 1024
+mlp_sizes = [480, 1024, 1024, 512, 256, 1]
+num_iters = 10
+
+class TestMLP(unittest.TestCase):
+
+    def test_creation(self):
+        MLP(mlp_sizes)
+
+    def test_numeric(self):
+        mlp = MLP(mlp_sizes).cuda()
+
+        mlp_layers = []
+        for i in range(mlp.num_layers):
+            linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1])
+            mlp.weights[i].data.copy_(linear.weight)
+            mlp.biases[i].data.copy_(linear.bias)
+            mlp_layers.append(linear)
+            mlp_layers.append(nn.ReLU(inplace=True))
+
+        ref_mlp = nn.Sequential(*mlp_layers).cuda()
+
+        test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.).requires_grad_()
+        ref_input = test_input.clone().detach().requires_grad_()
+        mlp_out = mlp(test_input)
+        ref_out = ref_mlp(ref_input)
+        np.testing.assert_allclose(
+            mlp_out.detach().cpu().numpy(),
+            ref_out.detach().cpu().numpy(),
+            atol=1e-7, rtol=1e-5)
+
+        # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out
+        mlp_out.mean().mul(10.).backward()
+        ref_out.mean().mul(10.).backward()
+        np.testing.assert_allclose(
+            test_input.grad.detach().cpu().numpy(),
+            ref_input.grad.detach().cpu().numpy(),
+            atol=0, rtol=1e-5)
+        np.testing.assert_allclose(
+            mlp.biases[0].grad.detach().cpu().numpy(),
+            ref_mlp[0].bias.grad.detach().cpu().numpy(),
+            atol=1e-7, rtol=1e-5)
+
+    def test_no_bias(self):
+        for use_activation in ['none', 'relu', 'sigmoid']:
+            mlp = MLP(mlp_sizes, bias=False, activation=use_activation).cuda()
+
+            mlp_layers = []
+            for i in range(mlp.num_layers):
+                linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=False)
+                mlp.weights[i].data.copy_(linear.weight)
+                mlp_layers.append(linear)
+                if use_activation == 'relu':
+                    mlp_layers.append(nn.ReLU(inplace=True))
+                if use_activation == 'sigmoid':
+                    mlp_layers.append(nn.Sigmoid())
+
+            ref_mlp = nn.Sequential(*mlp_layers).cuda()
+
+            test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.).requires_grad_()
+            ref_input = test_input.clone().detach().requires_grad_()
+            mlp_out = mlp(test_input)
+            ref_out = ref_mlp(ref_input)
+            np.testing.assert_allclose(
+                mlp_out.detach().cpu().numpy(),
+                ref_out.detach().cpu().numpy(),
+                atol=1e-7, rtol=1e-5)
+
+            # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out
+            mlp_out.mean().mul(10.).backward()
+            ref_out.mean().mul(10.).backward()
+            np.testing.assert_allclose(
+                test_input.grad.detach().cpu().numpy(),
+                ref_input.grad.detach().cpu().numpy(),
+                atol=0, rtol=100)
+            np.testing.assert_allclose(
+                mlp.weights[0].grad.detach().cpu().numpy(),
+                ref_mlp[0].weight.grad.detach().cpu().numpy(),
+                atol=1e-7, rtol=100)
+
+    def test_with_bias(self):
+        for use_activation in ['none', 'relu', 'sigmoid']:
+            mlp = MLP(mlp_sizes, bias=True, activation=use_activation).cuda()
+
+            mlp_layers = []
+            for i in range(mlp.num_layers):
+                linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1], bias=True)
+                mlp.weights[i].data.copy_(linear.weight)
+                mlp.biases[i].data.copy_(linear.bias)
+                mlp_layers.append(linear)
+                if use_activation == 'relu':
+                    mlp_layers.append(nn.ReLU(inplace=True))
+                if use_activation == 'sigmoid':
+                    mlp_layers.append(nn.Sigmoid())
+
+            ref_mlp = nn.Sequential(*mlp_layers).cuda()
+
+            test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.).requires_grad_()
+            ref_input = test_input.clone().detach().requires_grad_()
+            mlp_out = mlp(test_input)
+            ref_out = ref_mlp(ref_input)
+            np.testing.assert_allclose(
+                mlp_out.detach().cpu().numpy(),
+                ref_out.detach().cpu().numpy(),
+                atol=1e-7, rtol=1e-5)
+
+            # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out
+            mlp_out.mean().mul(10.).backward()
+            ref_out.mean().mul(10.).backward()
+            np.testing.assert_allclose(
+                test_input.grad.detach().cpu().numpy(),
+                ref_input.grad.detach().cpu().numpy(),
+                atol=0, rtol=1)
+            np.testing.assert_allclose(
+                mlp.weights[0].grad.detach().cpu().numpy(),
+                ref_mlp[0].weight.grad.detach().cpu().numpy(),
+                atol=1e-7, rtol=1)
+            np.testing.assert_allclose(
+                mlp.biases[0].grad.detach().cpu().numpy(),
+                ref_mlp[0].bias.grad.detach().cpu().numpy(),
+                atol=1e-7, rtol=1e-5)
+
+    def test_no_grad(self):
+        mlp = MLP(mlp_sizes).cuda()
+
+        mlp_layers = []
+        for i in range(mlp.num_layers):
+            linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1])
+            mlp.weights[i].data.copy_(linear.weight)
+            mlp.biases[i].data.copy_(linear.bias)
+            mlp_layers.append(linear)
+            mlp_layers.append(nn.ReLU(inplace=True))
+
+        ref_mlp = nn.Sequential(*mlp_layers).cuda()
+
+        test_input = torch.empty(batch_size, mlp_sizes[0], device="cuda").uniform_(-1., 1.)
+        ref_input = test_input.clone().detach()
+        mlp_out = mlp(test_input)
+        ref_out = ref_mlp(ref_input)
+        np.testing.assert_allclose(
+            mlp_out.detach().cpu().numpy(),
+            ref_out.detach().cpu().numpy(),
+            atol=1e-7, rtol=1e-5)
+
+        # Use mean value as scalar loss. Multiply 10 to make it big enough not zero out
+        mlp_out.mean().mul(10.).backward()
+        ref_out.mean().mul(10.).backward()
+        np.testing.assert_allclose(
+            mlp.weights[0].grad.detach().cpu().numpy(),
+            ref_mlp[0].weight.grad.detach().cpu().numpy(),
+            atol=1e-7, rtol=1e-5)
+
+
+    def test_performance_half(self):
+        mlp = MLP(mlp_sizes).cuda().half()
+
+        mlp_layers = []
+        for i in range(mlp.num_layers):
+            linear = nn.Linear(mlp_sizes[i], mlp_sizes[i + 1])
+            mlp.weights[i].data.copy_(linear.weight)
+            mlp.biases[i].data.copy_(linear.bias)
+            mlp_layers.append(linear)
+            mlp_layers.append(nn.ReLU(inplace=True))
+
+        ref_mlp = nn.Sequential(*mlp_layers).cuda().half()
+
+        test_input = torch.empty(
+            batch_size, mlp_sizes[0], device="cuda", dtype=torch.half).fill_(10.).requires_grad_()
+        ref_input = torch.empty(
+            batch_size, mlp_sizes[0], device="cuda", dtype=torch.half).fill_(10.).requires_grad_()
+
+        # Warm up GPU
+        for _ in range(100):
+            ref_out = ref_mlp(ref_input)
+            ref_loss = ref_out.mean()
+            ref_mlp.zero_grad()
+            ref_loss.backward()
+            mlp_out = mlp(test_input)
+            test_loss = mlp_out.mean()
+            mlp.zero_grad()
+            test_loss.backward()
+
+        torch.cuda.profiler.start()
+        torch.cuda.synchronize()
+        start_time = time()
+        for _ in range(num_iters):
+            ref_out = ref_mlp(ref_input)
+            ref_loss = ref_out.mean()
+            ref_mlp.zero_grad()
+            ref_loss.backward()
+        torch.cuda.synchronize()
+        stop_time = time()
+        print(F"\nPytorch MLP time {(stop_time - start_time) * 1000. / num_iters:.4f} ms")
+
+        torch.cuda.synchronize()
+        start_time = time()
+        for _ in range(num_iters):
+            mlp_out = mlp(test_input)
+            test_loss = mlp_out.mean()
+            mlp.zero_grad()
+            test_loss.backward()
+        torch.cuda.synchronize()
+        stop_time = time()
+        print(F"C++ MLP time {(stop_time - start_time) * 1000. / num_iters:.4f} ms")
+        torch.cuda.profiler.stop()
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_dist_adam.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_dist_adam.py
new file mode 100644
index 0000000000..c44180d3fc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_dist_adam.py
@@ -0,0 +1,183 @@
+import argparse
+import random
+import sys
+
+import torch
+from torch.nn.parallel import DistributedDataParallel as DDP
+
+from apex import amp
+from apex.optimizers import FusedAdam
+from apex.contrib.optimizers.distributed_fused_adam import DistributedFusedAdam
+
+
+class TestModel(torch.nn.Module):
+    def __init__(self, args):
+        super(TestModel, self).__init__()
+
+        self.linear = torch.nn.Sequential(*[torch.nn.Linear(args.dim, args.dim, bias=args.bias) for _ in range(args.layers)])
+
+    def forward(self, x):
+        return self.linear(x)
+
+def setup(args):
+    ## Model
+    ref_model = TestModel(args).cuda()
+    dist_model = TestModel(args).cuda()
+
+    # Same weights
+    with torch.no_grad():
+        for dp, rp in zip(dist_model.parameters(), ref_model.parameters()):
+            dp.data.copy_(rp.data)
+
+    dist_model = dist_model.half()
+
+
+    ## Optimizer
+    # same hyperparameters
+    ref_opt_args = { 'lr': 1e-3, 'eps': 1e-6, 'weight_decay': 0.01 }
+    ref_opt = FusedAdam(ref_model.parameters(), **ref_opt_args)
+
+    dist_opt_args = ref_opt_args.copy()
+    dist_opt_args.update( {'overlap_reductions' : False} )
+    dist_opt_args.update( {'process_group_size' : args.n_gpu} )
+    dist_opt_args.update( {'dwu_group_size' : args.dwu_group_size} )
+    dist_opt_args.update( {'dwu_num_blocks' : 1} )
+    dist_opt_args.update( {'dwu_num_chunks' : 1} )
+    dist_opt = DistributedFusedAdam(dist_model.parameters(), **dist_opt_args)
+    dist_opt.set_global_scale(1.)
+    
+    ## amp-init
+    amp_args = { 'loss_scale' : 'dynamic' , 'opt_level' : 'O2'}
+    ref_model, ref_opt = amp.initialize(ref_model, ref_opt, **amp_args)
+    
+   
+    ## DDP
+    ref_model = DDP(ref_model, device_ids=[args.rank])
+    with torch.no_grad():
+        for dp in dist_model.parameters():
+             torch.distributed.broadcast(dp.data, src=0)
+        for rp in ref_model.parameters():
+            torch.distributed.broadcast(rp.data, src=0)
+    torch.cuda.synchronize()
+    torch.distributed.barrier()
+    if get_rank() == 0:
+        print(f'dist opt with {args.n_gpu} GPUs')
+
+    return ref_model, ref_opt, dist_model, dist_opt
+
+def parse_args():
+    parser = argparse.ArgumentParser()
+
+    parser.add_argument('--local_rank', type=int, default=-1)
+    parser.add_argument('--steps', type=int, default=20)
+    parser.add_argument('--batch', type=int, default=32)
+    parser.add_argument('--dim', type=int, default=4)
+    parser.add_argument('--layers', type=int, default=2)
+    parser.add_argument('--bias', action='store_true')
+    parser.add_argument('--atol', type=float, default=1e-3)
+    parser.add_argument('--rtol', type=float, default=1)
+    parser.add_argument('--dwu_group_size', type=float, default=1)
+
+    args = parser.parse_args()
+
+    return args
+
+def setup_env(args):
+    torch.cuda.set_device(args.local_rank)
+    torch.distributed.init_process_group(backend='nccl', init_method='env://')
+    args.rank = torch.distributed.get_rank()
+    args.n_gpu = torch.distributed.get_world_size()
+
+    seed = 42 + get_rank()
+
+    random.seed(seed)
+    torch.manual_seed(seed)
+
+    return args
+
+def get_rank():
+    return torch.distributed.get_rank()
+
+def main():
+    args = parse_args()
+    args = setup_env(args)
+    tol_args = { 'atol' : args.atol, 'rtol' : args.rtol }
+
+    torch.set_printoptions(precision=16)
+
+    ref_model, ref_opt, dist_model, dist_opt = setup(args)
+
+    # lazy_init not called yet, initialize stash
+    stash = ref_opt._amp_stash
+    stash.all_fp16_params, stash.all_fp32_from_fp16_params = [], []
+
+    # make sure everything from _first_step_init_ is ready before training
+    # e.g. registering allreduce_hook
+    # so that gradients are copied/reduced when necessary
+    dist_opt._init_everything()
+
+    for i in range(args.steps):
+        x_ref = torch.randn(args.batch, args.dim, dtype=torch.half).cuda().requires_grad_(True)
+        x_dist = x_ref.clone().detach().requires_grad_(True)
+        
+        if get_rank() == 0:
+            print(f'[{i}] Checking input')
+        #print("x_ref:", x_ref.flatten()[:10])
+        #print("x_dist:", x_dist.flatten()[:10])
+        assert(torch.allclose(x_ref, x_dist, **tol_args))
+
+        y_ref = ref_model(x_ref).half()
+        y_dist = dist_model(x_dist)
+
+        if get_rank() == 0:
+            print(f'[{i}] Checking output')
+        #print("y_ref:", y_ref.flatten()[:10])
+        #print("y_dist:", y_dist.flatten()[:10])
+        assert(torch.allclose(y_ref, y_dist, **tol_args))
+
+        dy = torch.randn_like(y_ref)
+
+        y_ref.backward(dy)
+        y_dist.backward(dy)
+
+        if get_rank() == 0:
+            print(f'[{i}] Checking gradients')
+        torch.distributed.barrier()
+        torch.cuda.synchronize()
+        assert(torch.allclose(x_ref.grad, x_dist.grad, **tol_args))
+
+        # gradient all-reduce within distributed optimizer
+        dist_opt.complete_reductions()
+
+        if get_rank() == 0:
+            print(f'[{i}] Stepping')
+        ref_opt.step()
+        dist_opt.step()
+
+        torch.cuda.synchronize()
+        torch.distributed.barrier()
+        print('Checking new weights')
+        if get_rank() == 0:
+            print("ref param:", ref_model.module.linear[0].weight)
+            print("dist param:", dist_model.linear[0].weight)
+        
+        for i, (rp, dp) in enumerate(zip(ref_model.parameters(), dist_model.parameters())):
+            if not torch.allclose(rp, dp, **tol_args):
+                if get_rank() == 0:
+                    print(f'Rank: {get_rank()}, Param: {i}')
+                    print(f'ref: {rp.sum().item()}, dist: {dp.sum().item()}')
+                    print(rp)
+                    print(dp)
+    
+                    print(torch.abs(rp-dp) > tol_args['atol'])
+                    sys.exit(0)
+
+        # zero grads
+        for rp, dp in zip(ref_model.parameters(), dist_model.parameters()):
+            rp.grad = None
+            dp.grad = None
+
+
+if __name__ == "__main__":
+    main()
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_fused_novograd.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_fused_novograd.py
new file mode 100644
index 0000000000..fa94e71029
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_fused_novograd.py
@@ -0,0 +1,170 @@
+import torch
+from torch.optim import Optimizer
+import math
+import apex
+import unittest
+
+from test_fused_optimizer import TestFusedOptimizer
+from itertools import product
+
+class Novograd(Optimizer):
+    """
+    Implements Novograd algorithm.
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.95, 0))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        grad_averaging: gradient averaging
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.95, 0), eps=1e-8,
+                 weight_decay=0, grad_averaging=False, amsgrad=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                      weight_decay=weight_decay,
+                      grad_averaging=grad_averaging,
+                      amsgrad=amsgrad)
+
+        super(Novograd, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Novograd, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+            and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Sparse gradients are not supported.')
+                amsgrad = group['amsgrad']
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros([]).to(state['exp_avg'].device)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros([]).to(state['exp_avg'].device)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                if amsgrad:
+                    max_exp_avg_sq = state['max_exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                norm = torch.sum(torch.pow(grad, 2))
+
+                if exp_avg_sq == 0:
+                    exp_avg_sq.copy_(norm)
+                else:
+                    exp_avg_sq.mul_(beta2).add_(norm, alpha=1 - beta2)
+
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group['eps'])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                grad.div_(denom)
+                if group['weight_decay'] != 0:
+                    grad.add_(p.data, alpha=group['weight_decay'])
+                if group['grad_averaging']:
+                    grad.mul_(1 - beta1)
+                exp_avg.mul_(beta1).add_(grad)
+
+                p.data.add_(exp_avg, alpha=-group['lr'])
+        
+        return loss
+
+
+class TestFusedNovoGrad(TestFusedOptimizer):
+
+    def __init__(self, *args, **kwargs):
+        super(TestFusedNovoGrad, self).__init__(*args, **kwargs)
+
+        # The options for NovoGrad and FusedNovoGrad are very specific if they
+        # are expected to behave the same.
+        self.options = {'lr':1e-3, 'betas':(0.95, 0), 'eps':1e-8,
+                 'weight_decay':0, 'grad_averaging':False, 'amsgrad':False}
+        
+        self.tst_options = {'lr':1e-3, 'betas':(0.95, 0), 'eps':1e-8,
+                 'weight_decay':0, 'grad_averaging':False, 'amsgrad':False, 
+                 'bias_correction':False, 'reg_inside_moment':True, 
+                 'norm_type':2, 'init_zero':False, 'set_grad_none':True}
+
+        self.ref_optim = Novograd
+        self.fused_optim = apex.optimizers.FusedNovoGrad
+
+    def test_float(self):
+        self.gen_single_type_test(param_type=torch.float)
+
+    def test_half(self):
+        self.gen_single_type_test(param_type=torch.float16)
+
+    @unittest.skipIf(torch.cuda.device_count()<2, "more than 1 GPU required")
+    def test_multi_device(self):
+        devices = ("cuda:1", "cuda:0")
+        for current_dev, tensor_dev in product(devices, devices):
+            with torch.cuda.device(current_dev):
+                torch.cuda.synchronize()
+                self.gen_single_type_test(param_type=torch.float, device=tensor_dev)
+                
+
+    def test_multi_params(self):
+        sizes = [[4096, 1024], [4096], [4096, 2048], [32320, 1024], [1]]
+
+        tensors = []
+        for size in sizes:
+            tensors.append(torch.rand(size, dtype=torch.float, device="cuda"))
+        ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(
+            tensors, self.options, self.tst_options
+        )
+
+        for _ in range(self.iters):
+            self.gen_grad(ref_param, tst_param)
+            ref_optim.step()
+            tst_optim.step()
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_fused_optimizer.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_fused_optimizer.py
new file mode 100644
index 0000000000..05ed85abc9
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_fused_optimizer.py
@@ -0,0 +1,284 @@
+import unittest
+import os
+import random
+
+import math
+import torch
+import apex
+from itertools import product
+from torch.optim import Optimizer
+
+class TestFusedOptimizer(unittest.TestCase):
+    def setUp(self, max_abs_diff=1e-3, max_rel_diff=1, iters=7):
+        self.max_abs_diff = max_abs_diff
+        self.max_rel_diff = max_rel_diff
+        self.iters = iters
+        torch.cuda.manual_seed(9876)
+
+    def tearDown(self):
+        pass
+
+    def gen_param_optim(self, tensors, options, tst_options=None):
+        
+        # Adding this to make backward compatible with existing tests. Just in
+        # case "tst_options" are not provided, it gets a copy of options
+        # which contains the parameters for the reference optimizer
+        if tst_options == None:
+            tst_options = options
+            
+        ref_param = []
+        tst_param = []
+        for tensor in tensors:
+            ref_param.append(torch.nn.Parameter(tensor.clone()))
+            tst_param.append(torch.nn.Parameter(tensor.clone()))
+
+        ref_optim = self.ref_optim(ref_param, **options)
+        tst_optim = self.fused_optim(tst_param, **tst_options)
+
+        return (ref_param, tst_param, ref_optim, tst_optim)
+
+    def gen_grad(self, ref_param, tst_param):
+        for p_ref, p_tst in zip(ref_param, tst_param):
+            p_ref.grad = torch.rand_like(p_ref)
+            p_tst.grad = p_ref.grad
+
+    def gen_mixed_grad(self, ref_param, tst_param, scale=1.0):
+        half_grads = []
+        for p_ref, p_tst in zip(ref_param, tst_param):
+            half_grads.append(torch.rand_like(p_ref).half())
+            p_ref.grad = half_grads[-1].float() / scale
+        return half_grads
+
+    def get_max_diff(self, ref_param, tst_param):
+        max_abs_diff = max_rel_diff = 0
+        for p_ref, p_tst in zip(ref_param, tst_param):
+            max_abs_diff_p = (p_ref - p_tst).abs().max().item()
+            max_rel_diff_p = ((p_ref - p_tst) / p_ref).abs().max().item()
+
+            if max_abs_diff_p > max_abs_diff:  max_abs_diff = max_abs_diff_p
+            if max_rel_diff_p > max_rel_diff:  max_rel_diff = max_rel_diff_p
+
+        return max_abs_diff, max_rel_diff
+
+    def gen_single_type_test(self, param_type=torch.float, device='cuda'):
+        nelem = 278011
+
+        # Some ref and test optimizers may require different set of options.
+        # This is a quick workaround to add that functionality while making 
+        # minimum changes in existing code.
+        # If there is no "tst_options" field provided, safe to initialize
+        # the test optimizer with the parameters of reference optimizer.
+        if not hasattr(self, 'tst_options'):
+            self.tst_options = self.options
+
+        tensor = torch.rand(nelem, dtype=param_type, device=device)
+
+        ref_param, tst_param, ref_optim, tst_optim = \
+            self.gen_param_optim([tensor], self.options, self.tst_options)
+
+        for i in range(self.iters):
+            self.gen_grad(ref_param, tst_param)
+            ref_optim.step()
+            tst_optim.step()
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+
+class TestFusedAdam(TestFusedOptimizer):
+
+    def __init__(self, *args, **kwargs):
+        super(TestFusedAdam, self).__init__(*args, **kwargs)
+        self.options = {'lr':5e-4, 'betas':(0.9, 0.999), 'eps':1e-08,
+            'weight_decay': 0, 'amsgrad': False}
+        self.ref_optim = torch.optim.Adam
+        self.fused_optim = apex.optimizers.FusedAdam
+
+    def test_float(self):
+        self.gen_single_type_test(param_type=torch.float)
+
+    def test_half(self):
+        self.gen_single_type_test(param_type=torch.float16)
+
+    @unittest.skipIf(torch.cuda.device_count()<2, "more than 1 GPU required")
+    def test_multi_device(self):
+        devices = ("cuda:0", "cuda:1")
+        for current_dev, tensor_dev in product(devices, devices):
+            with torch.cuda.device(current_dev):
+                self.gen_single_type_test(param_type=torch.float, device=tensor_dev)
+
+    @unittest.skip('Disable until 8/1/2019 adam/adamw upstream picked')
+    def test_multi_params(self):
+        sizes = [[4096, 1024], [4096], [4096, 2048], [32320, 1024], [1]]
+
+        tensors = []
+        for size in sizes:
+            tensors.append(torch.rand(size, dtype=torch.float, device='cuda'))
+        ref_param, tst_param, ref_optim, tst_optim = \
+            self.gen_param_optim(tensors, self.options)
+
+        for i in range(self.iters):
+            self.gen_grad(ref_param, tst_param)
+            ref_optim.step()
+            tst_optim.step()
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+    @unittest.skip('No longer support fuse scaling')
+    def test_scale(self):
+        nelem = 278011
+        tensor = torch.rand(nelem, dtype=torch.float, device='cuda')
+        ref_param, tst_param, ref_optim, tst_optim = \
+            self.gen_param_optim([tensor], self.options)
+
+        for i in range(self.iters):
+            scale = random.random() * 1000
+            half_grads = self.gen_mixed_grad(ref_param, tst_param, scale)
+            ref_optim.step()
+            tst_optim.step(grads=half_grads, scale=scale)
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+    @unittest.skip('No longer support output fp16 param')
+    def test_fp16_output(self):
+        nelem = 278011
+
+        tensor = torch.rand(nelem, dtype=torch.float, device='cuda')
+        ref_param, tst_param, ref_optim, tst_optim = \
+            self.gen_param_optim([tensor], self.options)
+
+        fp16_param = torch.nn.Parameter(tensor.clone().half())
+
+        for i in range(self.iters):
+            half_grads = self.gen_mixed_grad(ref_param, tst_param)
+            ref_optim.step()
+            tst_optim.step(grads=half_grads, output_params=[fp16_param])
+
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+            max_abs_diff, max_rel_diff = self.get_max_diff(tst_param, \
+                [fp16_param.float()])
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+    def test_adam_option(self):
+        nelem = 1
+        adam_option = {'lr':0.01, 'betas':(0.6, 0.9), 'eps':3e-06,
+            'weight_decay':0, 'amsgrad':False}
+
+        tensor = torch.rand(nelem, dtype=torch.float, device='cuda')
+        ref_param, tst_param, ref_optim, tst_optim = \
+            self.gen_param_optim([tensor], adam_option)
+
+        for i in range(self.iters):
+            self.gen_grad(ref_param, tst_param)
+            ref_optim.step()
+            tst_optim.step()
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+
+class TestFusedAdagrad(TestFusedOptimizer):
+    def __init__(self, *args, **kwargs):
+        super(TestFusedAdagrad, self).__init__(*args, **kwargs)
+        self.options = {"lr": 5e-4, "eps": 1e-08, "weight_decay": 1.0e-5}
+        self.ref_optim = torch.optim.Adagrad
+        self.fused_optim = apex.optimizers.FusedAdagrad
+
+    def test_float(self):
+        self.gen_single_type_test(param_type=torch.float)
+
+    @unittest.skip("PyTorch optimizer is not numerically correct for fp16")
+    def test_half(self):
+        self.gen_single_type_test(param_type=torch.float16)
+
+    @unittest.skipIf(torch.cuda.device_count()<2, "more than 1 GPU required")
+    def test_multi_device(self):
+        devices = ("cuda:0", "cuda:1")
+        for current_dev, tensor_dev in product(devices, devices):
+            with torch.cuda.device(current_dev):
+                self.gen_single_type_test(param_type=torch.float, device=tensor_dev)
+
+
+    def test_multi_params(self):
+        sizes = [[4096, 1024], [4096], [4096, 2048], [32320, 1024], [1]]
+        adagrad_option = {"lr": 5e-4, "eps": 1e-08, "weight_decay": 0}
+
+        tensors = []
+        for size in sizes:
+            tensors.append(torch.rand(size, dtype=torch.float, device="cuda"))
+        ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(
+            tensors, adagrad_option
+        )
+
+        for _ in range(self.iters):
+            self.gen_grad(ref_param, tst_param)
+            ref_optim.step()
+            tst_optim.step()
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+    @unittest.skipIf(torch.cuda.device_count()<2, "more than 1 GPU required")
+    def test_multi_params_different_devices_throws(self):
+        sizes = [[4096, 1024], [4096], [4096, 2048], [32320, 1024], [1]]
+        adagrad_option = {"lr": 5e-4, "eps": 1e-08, "weight_decay": 0}
+
+        tensors = []
+        for i, size in enumerate(sizes):
+            tensors.append(torch.rand(size, dtype=torch.float, device="cuda:"+str(i % 2)))
+        ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(
+            tensors, adagrad_option
+        )
+        self.gen_grad(ref_param, tst_param)
+        with self.assertRaisesRegex(RuntimeError, "not on the same device"):
+            tst_optim.step()
+
+    def test_adagrad_option(self):
+        nelem = 1
+        adagrad_option = {"lr": 0.01, "eps": 3e-06, "weight_decay": 0}
+
+        tensor = torch.rand(nelem, dtype=torch.float, device="cuda")
+        ref_param, tst_param, ref_optim, tst_optim = self.gen_param_optim(
+            [tensor], adagrad_option
+        )
+
+        for _ in range(self.iters):
+            self.gen_grad(ref_param, tst_param)
+            ref_optim.step()
+            tst_optim.step()
+            max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+
+            self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+            self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+
+class TestFusedSGD(TestFusedOptimizer):
+    def __init__(self, *args, **kwargs):
+        super(TestFusedSGD, self).__init__(*args, **kwargs)
+        self.options = {"lr": .25, "momentum": .125}
+        self.ref_optim = torch.optim.SGD
+        self.fused_optim = apex.optimizers.FusedSGD
+
+    def test_float(self):
+        self.gen_single_type_test(param_type=torch.float)
+
+    def test_half(self):
+        self.gen_single_type_test(param_type=torch.float16)
+
+    @unittest.skipIf(torch.cuda.device_count()<2, "more than 1 GPU required")
+    def test_multi_device(self):
+        devices = ("cuda:0", "cuda:1")
+        for current_dev, tensor_dev in product(devices, devices):
+            with torch.cuda.device(current_dev):
+                self.gen_single_type_test(param_type=torch.float, device=tensor_dev)
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_lamb.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_lamb.py
new file mode 100644
index 0000000000..e7186c1b0b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_optimizers/test_lamb.py
@@ -0,0 +1,270 @@
+import unittest
+import os
+
+import torch
+from torch.optim import Optimizer
+import apex
+from apex.multi_tensor_apply import multi_tensor_applier
+from itertools import product
+
+class RefLAMB(Optimizer):
+    r"""Implements Lamb algorithm.
+
+    It has been proposed in `Large Batch Optimization for Deep Learning: Training BERT in 76 minutes`_.
+
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-6)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0.01)
+
+    .. _Large Batch Optimization for Deep Learning: Training BERT in 76 minutes:
+        https://arxiv.org/abs/1904.00962
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-6, weight_decay=0.01):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps, weight_decay=weight_decay)
+        super(RefLAMB, self).__init__(params, defaults)
+        if multi_tensor_applier.available:
+            import amp_C
+            self.multi_tensor_l2norm=amp_C.multi_tensor_l2norm
+            # Skip buffer
+            self._dummy_overflow_buf = torch.tensor([0], dtype=torch.int, device=self.param_groups[0]["params"][0].device)
+            self.multi_tensor_lamb = amp_C.multi_tensor_lamb
+        else:
+            raise RuntimeError('apex.optimizers.FusedLAMB requires cuda extensions')
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        # create separate grad lists for fp32 and fp16 params
+        g_all_32, g_all_16 = [], []
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                if p.dtype == torch.float32:
+                    g_all_32.append(p.grad.data)
+                elif p.dtype == torch.float16:
+                    g_all_16.append(p.grad.data)
+                else:
+                    raise RuntimeError('FusedLAMB only support fp16 and fp32.')
+
+        device = self.param_groups[0]["params"][0].device
+        g_norm_32, g_norm_16 = torch.zeros(1, device=device), torch.zeros(1, device=device)
+        # compute grad norm for two lists
+        if len(g_all_32) > 0:
+            g_norm_32 = multi_tensor_applier(self.multi_tensor_l2norm,
+                                             self._dummy_overflow_buf,
+                                             [g_all_32], False)[0]
+        if len(g_all_16) > 0:
+            g_norm_16 = multi_tensor_applier(self.multi_tensor_l2norm,
+                                             self._dummy_overflow_buf,
+                                             [g_all_16], False)[0]
+
+        # blend two grad norms to get global grad norm
+        global_grad_norm = multi_tensor_applier(self.multi_tensor_l2norm,
+                                                self._dummy_overflow_buf,
+                                                [[g_norm_32, g_norm_16]],
+                                                False)[0]
+
+        max_grad_norm = 1.0
+        clipped_ratio = max_grad_norm / max(global_grad_norm, max_grad_norm)
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                p.grad.data *= clipped_ratio
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Lamb does not support sparse gradients, consider SparseAdam instad.')
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['m'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['v'] = torch.zeros_like(p.data)
+
+                m_t, v_t = state['m'], state['v']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                # m_t = beta1 * m + (1 - beta1) * g_t
+                m_t.mul_(beta1).add_(grad, alpha=1-beta1)
+                # v_t = beta2 * v + (1 - beta2) * (g_t * g_t)
+                v_t.mul_(beta2).addcmul_(grad, grad, value=1-beta2)
+
+                # Debiasing
+                m_t_hat = m_t / (1.0 - beta1 ** state['step'])
+                v_t_hat = v_t / (1.0 - beta2 ** state['step'])
+
+                update = m_t_hat / v_t_hat.sqrt().add(group['eps'])
+
+                if group['weight_decay'] != 0:
+                    update.add_(p.data, alpha=group['weight_decay'])
+
+                trust_ratio = 1.0
+                w_norm = p.data.pow(2).sum().sqrt()
+                g_norm = update.pow(2).sum().sqrt()
+                if w_norm > 0 and g_norm > 0:
+                    trust_ratio = w_norm / g_norm
+
+                state['w_norm'] = w_norm
+                state['g_norm'] = g_norm
+                state['trust_ratio'] = trust_ratio
+
+                step_size = group['lr']
+
+                p.data.add_(update, alpha=-step_size*trust_ratio)
+
+        return loss
+
+
+class TestFusedLAMB(unittest.TestCase):
+    def setUp(self, max_abs_diff=1e-3, max_rel_diff=1, iters=7):
+        self.max_abs_diff = max_abs_diff
+        self.max_rel_diff = max_rel_diff
+        self.iters = iters
+        torch.cuda.manual_seed(9876)
+
+    def tearDown(self):
+        pass
+
+    def gen_param_optim(self, tensors, lamb_option):
+        ref_param = []
+        tst_param = []
+        for tensor in tensors:
+            ref_param.append(torch.nn.Parameter(tensor.clone()))
+            tst_param.append(torch.nn.Parameter(tensor.clone()))
+
+        ref_optim = RefLAMB(ref_param, **lamb_option)
+        tst_optim = apex.optimizers.FusedLAMB(tst_param, use_nvlamb=True, **lamb_option)
+
+        return (ref_param, tst_param, ref_optim, tst_optim)
+
+    def gen_grad(self, ref_param, tst_param):
+        for p_ref, p_tst in zip(ref_param, tst_param):
+            p_ref.grad = torch.rand_like(p_ref)
+            p_tst.grad = p_ref.grad
+
+    def gen_mixed_grad(self, ref_param, tst_param, scale=1.0):
+        half_grads = []
+        for p_ref, _ in zip(ref_param, tst_param):
+            half_grads.append(torch.rand_like(p_ref).half())
+            p_ref.grad = half_grads[-1].float() / scale
+        return half_grads
+
+    def get_max_diff(self, ref_param, tst_param):
+        max_abs_diff = max_rel_diff = 0
+        for p_ref, p_tst in zip(ref_param, tst_param):
+            max_abs_diff_p = (p_ref - p_tst).abs().max().item()
+            max_rel_diff_p = ((p_ref - p_tst) / p_ref).abs().max().item()
+
+            if max_abs_diff_p > max_abs_diff:  max_abs_diff = max_abs_diff_p
+            if max_rel_diff_p > max_rel_diff:  max_rel_diff = max_rel_diff_p
+
+        return max_abs_diff, max_rel_diff
+
+    def gen_single_type_test(self, param_type=torch.float, device="cuda"):
+        nelem = 278011
+        tensor = torch.rand(nelem, dtype=param_type, device=device)
+        weight_decay = [0, 0.01]
+
+        for wd in weight_decay:
+            lamb_option = {'lr':5e-4, 'betas':(0.9, 0.999), 'eps':1e-08, 'weight_decay':wd}
+            ref_param, tst_param, ref_optim, tst_optim = \
+                self.gen_param_optim([tensor], lamb_option)
+
+            for i in range(self.iters):
+                self.gen_grad(ref_param, tst_param)
+                ref_optim.step()
+                torch.cuda.synchronize()
+                tst_optim.step()
+                torch.cuda.synchronize()
+                max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+
+                self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+                self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+    def test_float(self):
+        self.gen_single_type_test(param_type=torch.float)
+
+    @unittest.skip("PyTorch optimizer is not numerically correct for fp16")
+    def test_half(self):
+        self.gen_single_type_test(param_type=torch.float16)
+
+    @unittest.skipIf(torch.cuda.device_count()<2, "more than 1 GPU required")
+    def test_multi_device(self):
+        devices = ("cuda:0", "cuda:1")
+        for current_dev, tensor_dev in product(devices, devices):
+            with torch.cuda.device(current_dev):
+                self.gen_single_type_test(param_type=torch.float, device=tensor_dev)
+
+    def test_multi_params(self):
+        sizes = [[4096, 1024], [4096], [4096, 2048], [32320, 1024], [1]]
+        weight_decay = [0, 0.01]
+
+        for wd in weight_decay:
+            lamb_option = {'lr':5e-4, 'betas':(0.9, 0.999), 'eps':1e-08, 'weight_decay':wd}
+            tensors = []
+            for size in sizes:
+                tensors.append(torch.rand(size, dtype=torch.float, device='cuda'))
+            ref_param, tst_param, ref_optim, tst_optim = \
+                self.gen_param_optim(tensors, lamb_option)
+
+            for i in range(self.iters):
+                self.gen_grad(ref_param, tst_param)
+                ref_optim.step()
+                tst_optim.step()
+                max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+                self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+                self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+    def test_lamb_option(self):
+        nelem = 1
+        tensor = torch.rand(nelem, dtype=torch.float, device='cuda')
+        weight_decay = [0, 0.01]
+
+        for wd in weight_decay:
+            lamb_option = {'lr':0.01, 'betas':(0.6, 0.9), 'eps':3e-06, 'weight_decay':wd}
+            ref_param, tst_param, ref_optim, tst_optim = \
+                self.gen_param_optim([tensor], lamb_option)
+
+            for i in range(self.iters):
+                self.gen_grad(ref_param, tst_param)
+                ref_optim.step()
+                tst_optim.step()
+                max_abs_diff, max_rel_diff = self.get_max_diff(ref_param, tst_param)
+
+                self.assertLessEqual(max_abs_diff, self.max_abs_diff)
+                self.assertLessEqual(max_rel_diff, self.max_rel_diff)
+
+
+if __name__ == '__main__':
+    script_path = os.path.dirname(os.path.realpath(__file__))
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_data/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_data/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_data/test_pyprof_data.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_data/test_pyprof_data.py
new file mode 100644
index 0000000000..266212a764
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_data/test_pyprof_data.py
@@ -0,0 +1,43 @@
+import inspect
+import unittest
+
+from apex.pyprof.prof.data import Data
+from apex.pyprof.prof.prof import foo
+
+
+class TestPyProfData(unittest.TestCase):
+
+	def __init__(self, testName):
+		super().__init__(testName)
+
+	def setUp(self):
+		pass
+
+	def tearDown(self):
+		pass
+
+	def test_data(self):
+		kernels = [
+			{'kShortName': 'elementwise_kernel', 'kDuration': 2848, 'layer': [], 'trace': [], 'reprMarkers': [], 'marker': ["{'mod': 'Tensor', 'op': 'float', 'args': [{'name': '', 'type': 'tensor', 'shape': (18, 104, 160), 'dtype': 'bool'}]}"], 'seqMarker': ['to, seq = 60471'], 'seqId': [60471], 'subSeqId': 0, 'altSeqId': [], 'dir': 'fprop', 'mod': ['Tensor'], 'op': ['float'], 'tid': 1431533376, 'device': 0, 'stream': 7, 'grid': (585, 1, 1), 'block': (512, 1, 1), 'kLongName': 'void at::native::elementwise_kernel<512, 1, void at::native::gpu_kernel_impl<void at::native::copy_kernel_impl<float, bool>(at::TensorIterator&)::{lambda(bool)#1}>(at::TensorIterator&, void at::native::copy_kernel_impl<float, bool>(at::TensorIterator&)::{lambda(bool)#1} const&)::{lambda(int)#1}>(int, void at::native::gpu_kernel_impl<void at::native::copy_kernel_impl<float, bool>(at::TensorIterator&)::{lambda(bool)#1}>(at::TensorIterator&, void at::native::copy_kernel_impl<float, bool>(at::TensorIterator&)::{lambda(bool)#1} const&)::{lambda(int)#1})'},
+			{'kShortName': 'elementwise_kernel', 'kDuration': 201182, 'layer': [], 'trace': [], 'reprMarkers': [], 'marker': ["{'mod': 'Tensor', 'op': 'clone', 'args': [{'name': '', 'type': 'tensor', 'shape': (18, 4, 416, 640), 'dtype': 'float32'}]}"], 'seqMarker': ['clone, seq = 60161'], 'seqId': [60161], 'subSeqId': 0, 'altSeqId': [], 'dir': 'fprop', 'mod': ['Tensor'], 'op': ['clone'], 'tid': 1431533376, 'device': 0, 'stream': 7, 'grid': (37440, 1, 1), 'block': (128, 1, 1), 'kLongName': 'void at::native::elementwise_kernel<128, 4, void at::native::gpu_kernel_impl<void at::native::copy_kernel_impl<float, float>(at::TensorIterator&)::{lambda(float)#1}>(at::TensorIterator&, void at::native::copy_kernel_impl<float, float>(at::TensorIterator&)::{lambda(float)#1} const&)::{lambda(int)#2}>(int, void at::native::gpu_kernel_impl<void at::native::copy_kernel_impl<float, float>(at::TensorIterator&)::{lambda(float)#1}>(at::TensorIterator&, void at::native::copy_kernel_impl<float, float>(at::TensorIterator&)::{lambda(float)#1} const&)::{lambda(int)#2})'},
+		]
+
+		for k in kernels:
+			d = Data(k)
+			mod = k['mod']
+			op = k['op']
+			xx = foo(mod, op, d)
+			d.setParams(xx.params())
+
+
+def run_tests(test_name):
+	dummy = TestPyProfData(test_name)
+	test_cases = list(filter(lambda x: 'test_' in x, map(lambda x: x[0], inspect.getmembers(dummy, predicate=inspect.ismethod))))
+	print(f'Running tests for {test_name}')
+	suite = unittest.TestSuite()
+	for test_case in test_cases:
+		suite.addTest(TestPyProfData(test_case))
+	unittest.TextTestRunner().run(suite)
+
+if __name__ == '__main__':
+	run_tests('test_data')
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_nvtx/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_nvtx/__init__.py
new file mode 100644
index 0000000000..7cd4c1062a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_nvtx/__init__.py
@@ -0,0 +1 @@
+import test_pyprof_nvtx.TestPyProfNvtx as TestPyProfNvtx
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_nvtx/test_pyprof_nvtx.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_nvtx/test_pyprof_nvtx.py
new file mode 100644
index 0000000000..6f2c8d1e26
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_pyprof_nvtx/test_pyprof_nvtx.py
@@ -0,0 +1,526 @@
+import inspect
+import os
+import torch
+import torch.nn.functional as F
+import unittest
+
+from apex import pyprof
+pyprof.nvtx.init()
+
+# TODO: add tests for:
+# F.bilinear, F.l1_loss, F.multilabel_soft_margin_loss, F.multi_margin_loss
+
+class TestPyProfNvtx(unittest.TestCase):
+
+    def __init__(self, testName, dtype=torch.float16):
+        super().__init__(testName) 
+        self.dtype = dtype
+
+    def setUp(self):
+        pass
+
+    def tearDown(self):
+        pass
+
+    def test_conv1d(self):
+        # Data and weight tensors
+        tensor1d_in_conv = torch.randn(32, 3, 224, device='cuda', dtype=self.dtype)
+        tensor1d_in_conv_grouped = torch.randn(32, 6, 224, device='cuda', dtype=self.dtype)
+        conv1d_filter = torch.randn(16, 3, 3, device='cuda', dtype=self.dtype)
+        conv1d_bias = torch.ones(16, device='cuda', dtype=self.dtype)
+        # Vanilla conv1d
+        conv1d_out_vanilla = F.conv1d(tensor1d_in_conv, conv1d_filter)
+        # conv1d with bias
+        conv1d_out_with_bias = F.conv1d(tensor1d_in_conv, conv1d_filter, bias=conv1d_bias)
+        # conv1d - stride > 1
+        conv1d_out_strided = F.conv1d(tensor1d_in_conv, conv1d_filter, stride=2)
+        # conv1d - dilation > 1
+        conv1d_out_dilated = F.conv1d(tensor1d_in_conv, conv1d_filter, dilation=2)
+        # conv1d - groups > 1
+        conv1d_out_grouped = F.conv1d(tensor1d_in_conv_grouped, conv1d_filter, groups=2)
+        # conv1d - padding with zeros
+        conv1d_out_padding_zeros = F.conv1d(tensor1d_in_conv, conv1d_filter, padding=6)
+    
+    def test_conv2d(self):
+        # Data and weight tensors
+        tensor2d_in_conv = torch.randn(32, 3, 224, 224, device='cuda', dtype=self.dtype)
+        tensor2d_in_conv_grouped = torch.randn(32, 6, 224, 224, device='cuda', dtype=self.dtype)
+        conv2d_filter = torch.randn(16, 3, 3, 3, device='cuda', dtype=self.dtype)
+        conv2d_bias = torch.ones(16, device='cuda', dtype=self.dtype)
+        # Vanilla conv2d
+        conv2d_out_vanilla = F.conv2d(tensor2d_in_conv, conv2d_filter)
+        # conv2d with bias
+        conv2d_with_bias = F.conv2d(tensor2d_in_conv, conv2d_filter, bias=conv2d_bias)
+        # conv2d - stride > 1
+        conv2d_out_strided = F.conv2d(tensor2d_in_conv, conv2d_filter, stride=2)
+        # conv2d - dilation > 1
+        conv2d_out_dilated = F.conv2d(tensor2d_in_conv, conv2d_filter, dilation=2)
+        # conv2d - groups > 1
+        conv2d_out_grouped = F.conv2d(tensor2d_in_conv_grouped, conv2d_filter, groups=2)
+        # conv2d - padding with zeros
+        conv2d_out_padding_zeros = F.conv2d(tensor2d_in_conv, conv2d_filter, padding=6)
+    
+    
+    def test_conv3d(self):
+        # Data and weight tensors
+        tensor3d_in_conv = torch.randn(32, 3, 16, 224, 224, device='cuda', dtype=self.dtype)
+        tensor3d_in_conv_grouped = torch.randn(32, 6, 16, 224, 224, device='cuda', dtype=self.dtype)
+        conv3d_filter = torch.randn(16, 3, 3, 3, 3, device='cuda', dtype=self.dtype)
+        conv3d_bias = torch.ones(16, device='cuda', dtype=self.dtype)
+        # Vanilla conv3d
+        conv3d_out_vanilla = F.conv3d(tensor3d_in_conv, conv3d_filter)
+        # conv3d - stride > 1
+        conv3d_out_strided = F.conv3d(tensor3d_in_conv, conv3d_filter, stride=2)
+        # conv3d - dilation > 1
+        conv3d_out_dilated = F.conv3d(tensor3d_in_conv, conv3d_filter, dilation=2)
+        # conv3d - groups > 1
+        conv3d_out_grouped = F.conv3d(tensor3d_in_conv_grouped, conv3d_filter, groups=2)
+        # conv3d - padding with zeros
+        conv3d_out_padding_zeros = F.conv3d(tensor3d_in_conv, conv3d_filter, padding=6)
+    
+    def test_conv_transpose1d(self):
+        # Data and weight tensors
+        conv_transpose1d_tensor = torch.randn(64, 16, 64, device='cuda', dtype=self.dtype)
+        conv_transpose1d_filter = torch.randn(16, 32, 3, device='cuda', dtype=self.dtype)
+        conv_transpose1d_bias = torch.randn(32, device='cuda', dtype=self.dtype)
+        # Conv transpose runs
+        conv_transpose1d_out = F.conv_transpose1d(conv_transpose1d_tensor, conv_transpose1d_filter)
+        conv_transpose1d_out_biased = F.conv_transpose1d(conv_transpose1d_tensor, conv_transpose1d_filter, bias=conv_transpose1d_bias)
+        conv_transpose1d_out_strided = F.conv_transpose1d(conv_transpose1d_tensor, conv_transpose1d_filter, stride=2)
+        conv_transpose1d_out_padded = F.conv_transpose1d(conv_transpose1d_tensor, conv_transpose1d_filter, padding=3)
+        conv_transpose1d_out2_padded = F.conv_transpose1d(conv_transpose1d_tensor, conv_transpose1d_filter, output_padding=2, dilation=3)
+        conv_transpose1d_out_grouped = F.conv_transpose1d(conv_transpose1d_tensor, conv_transpose1d_filter, groups=2)
+        conv_transpose1d_out_dilated = F.conv_transpose1d(conv_transpose1d_tensor, conv_transpose1d_filter, dilation=2)
+    
+    
+    def test_conv_transpose2d(self):
+        # Data and weight tensors
+        conv_transpose2d_tensor = torch.randn(64, 8, 5, 5, device='cuda', dtype=self.dtype)
+        conv_transpose2d_filter = torch.randn(8, 16, 3, 3, device='cuda', dtype=self.dtype)
+        conv_transpose2d_bias = torch.randn(16, device='cuda', dtype=self.dtype)
+        # Conv transpose runs
+        conv_transpose2d_out = F.conv_transpose2d(conv_transpose2d_tensor, conv_transpose2d_filter)
+        conv_transpose2d_out_biased = F.conv_transpose2d(conv_transpose2d_tensor, conv_transpose2d_filter, bias=conv_transpose2d_bias)
+        conv_transpose2d_out_strided = F.conv_transpose2d(conv_transpose2d_tensor, conv_transpose2d_filter, stride=2)
+        conv_transpose2d_out_padded = F.conv_transpose2d(conv_transpose2d_tensor, conv_transpose2d_filter, padding=3)
+        conv_transpose2d_out2_padded = F.conv_transpose2d(conv_transpose2d_tensor, conv_transpose2d_filter, output_padding=2, dilation=3)
+        conv_transpose2d_out_grouped = F.conv_transpose2d(conv_transpose2d_tensor, conv_transpose2d_filter, groups=2)
+        conv_transpose2d_out_dilated = F.conv_transpose2d(conv_transpose2d_tensor, conv_transpose2d_filter, dilation=2)
+    
+    def test_conv_transpose3d(self):
+        # Data and weight tensors
+        conv_transpose3d_tensor = torch.randn(20, 16, 50, 10, 20, device='cuda', dtype=self.dtype)
+        conv_transpose3d_filter = torch.randn(16, 33, 3, 3, 3, device='cuda', dtype=self.dtype)
+        conv_transpose3d_bias = torch.randn(33, device='cuda', dtype=self.dtype)
+        # Conv transpose runs
+        conv_transpose3d_out = F.conv_transpose3d(conv_transpose3d_tensor, conv_transpose3d_filter)
+        conv_transpose3d_out_biased = F.conv_transpose3d(conv_transpose3d_tensor, conv_transpose3d_filter, bias=conv_transpose3d_bias)
+        conv_transpose3d_out_strided = F.conv_transpose3d(conv_transpose3d_tensor, conv_transpose3d_filter, stride=2)
+        conv_transpose3d_out_padded = F.conv_transpose3d(conv_transpose3d_tensor, conv_transpose3d_filter, padding=3)
+        conv_transpose3d_out2_padded = F.conv_transpose3d(conv_transpose3d_tensor, conv_transpose3d_filter, output_padding=2, dilation=3)
+        conv_transpose3d_out_grouped = F.conv_transpose3d(conv_transpose3d_tensor, conv_transpose3d_filter, groups=2)
+        conv_transpose3d_out_dilated = F.conv_transpose3d(conv_transpose3d_tensor, conv_transpose3d_filter, dilation=2)
+    
+    def test_unfold(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        kernel_size = (4, 5)
+        inp_unf_dilated = F.unfold(inp, kernel_size, dilation=2)
+        inp_unf_padded = F.unfold(inp, kernel_size, padding=2)
+        inp_unf_strided = F.unfold(inp, kernel_size, stride=2)
+    
+    def test_fold(self):
+        inp = torch.randn(3, 20, 20, device='cuda', dtype=self.dtype)
+        inp_folded = F.fold(inp, (4, 5), (1, 1))
+    
+    def test_avg_pool1d(self):
+        inp = torch.randn(1, 1, 28, device='cuda', dtype=self.dtype)
+        out = F.avg_pool1d(inp, kernel_size=5, stride=2, padding=2, ceil_mode=True, count_include_pad=False)
+    
+    def test_avg_pool2d(self):
+        inp = torch.randn(1, 3, 224, 224, device='cuda', dtype=self.dtype)
+        out = F.avg_pool2d(inp, kernel_size=5, stride=2, padding=2, ceil_mode=True, count_include_pad=False)
+    
+    def test_avg_pool3d(self):
+        inp = torch.randn(1, 3, 16, 224, 224, device='cuda', dtype=self.dtype)
+        out = F.avg_pool3d(inp, kernel_size=5, stride=2, padding=2, ceil_mode=True, count_include_pad=False)
+    
+    def test_adaptive_avg_pool1d(self):
+        inp = torch.randn(1, 1, 28, device='cuda', dtype=self.dtype)
+        out = F.adaptive_avg_pool1d(inp, output_size=5) 
+    
+    def test_adaptive_avg_pool2d(self):
+        inp = torch.randn(1, 16, 32, 32, device='cuda', dtype=self.dtype)
+        out = F.adaptive_avg_pool2d(inp, output_size=5) 
+    
+    def test_adaptive_avg_pool3d(self):
+        inp = torch.randn(1, 16, 16, 32, 32, device='cuda', dtype=self.dtype)
+        out = F.adaptive_avg_pool3d(inp, output_size=5) 
+    
+    def test_max_pool1d(self):
+        inp = torch.randn(1, 16, 32, device='cuda', dtype=self.dtype)
+        out = F.max_pool1d(inp, kernel_size=5, stride=2, padding=2, return_indices=True, ceil_mode=True)
+    
+    def test_max_pool2d(self):
+        inp = torch.randn(1, 16, 32, 32, device='cuda', dtype=self.dtype)
+        out = F.max_pool2d(inp, kernel_size=5, stride=2, padding=2, return_indices=True, ceil_mode=True)
+    
+    def test_max_pool3d(self):
+        inp = torch.randn(1, 16, 16, 32, 32, device='cuda', dtype=self.dtype)
+        out = F.max_pool3d(inp, kernel_size=5, stride=2, padding=2, return_indices=True, ceil_mode=True)
+    
+    def test_adaptive_max_pool1d(self):
+        inp = torch.randn(1, 16, 28, device='cuda', dtype=self.dtype)
+        out = F.adaptive_max_pool1d(inp, output_size=5, return_indices=True) 
+    
+    def test_adaptive_max_pool2d(self):
+        inp = torch.randn(1, 16, 32, 32, device='cuda', dtype=self.dtype)
+        out = F.adaptive_max_pool2d(inp, output_size=5, return_indices=True) 
+    
+    def test_adaptive_max_pool3d(self):
+        inp = torch.randn(1, 16, 16, 32, 32, device='cuda', dtype=self.dtype)
+        out = F.adaptive_max_pool3d(inp, output_size=5, return_indices=True) 
+    
+    def test_max_unpool1d(self):
+        inp = torch.randn(1, 16, 32, device='cuda', dtype=self.dtype)
+        output, indices = F.max_pool1d(inp, kernel_size=5, stride=2, padding=2, return_indices=True, ceil_mode=True)
+        output = F.max_unpool1d(output, indices, kernel_size=2, stride=2, padding=2)
+    
+    def test_max_unpool2d(self):
+        inp = torch.randn(1, 16, 32, 32, device='cuda', dtype=self.dtype)
+        output, indices = F.max_pool2d(inp, kernel_size=5, stride=2, padding=2, return_indices=True, ceil_mode=True)
+        output = F.max_unpool2d(output, indices, kernel_size=2, stride=2, padding=2)
+    
+    def test_max_unpool3d(self):
+        inp = torch.randn(1, 16, 8, 32, 32, device='cuda', dtype=self.dtype)
+        output, indices = F.max_pool3d(inp, kernel_size=5, stride=2, padding=2, return_indices=True, ceil_mode=True)
+        output = F.max_unpool3d(output, indices, kernel_size=2, stride=2, padding=2)
+    
+    def test_lp_pool1d(self):
+        inp = torch.randn(1, 32, 64, device='cuda', dtype=self.dtype)
+        output = F.lp_pool1d(inp, 2, 3, stride=2, ceil_mode=True)
+    
+    def test_lp_pool2d(self):
+        #torch.nn.LPPool2d(norm_type, kernel_size, stride=None, ceil_mode=False)
+        inp = torch.randn(1, 32, 64, 64, device='cuda', dtype=self.dtype)
+        output = F.lp_pool2d(inp, 2, 3, stride=2, ceil_mode=True)
+    
+    def test_threshold(self):
+        inp = torch.randn(1, 8, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.threshold(inp, 6, 6, inplace=False)
+    
+    def test_threshold_(self):
+        inp = torch.randn(1, 8, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.threshold_(inp, 6, 6)
+    
+    def test_relu(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.relu(inp, inplace=False)
+    
+    def test_relu_(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.relu_(inp)
+    
+    def test_hardtanh(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.hardtanh(inp, min_val=-1., max_val=1., inplace=False)
+    
+    def test_hardtanh_(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.hardtanh_(inp, min_val=-1., max_val=1.)
+    
+    def test_relu6(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.relu6(inp, inplace=False)
+    
+    def test_elu(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.elu(inp, alpha=1.0, inplace=False)
+    
+    def test_elu_(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.elu_(inp, alpha=1.0)
+    
+    def test_selu(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.selu(inp)
+    
+    def test_celu(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.celu(inp, alpha=1.0, inplace=False)
+    
+    def test_leaky_relu(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.leaky_relu(inp, negative_slope=0.01, inplace=False)
+    
+    def test_leaky_relu_(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.leaky_relu_(inp, negative_slope=0.01)
+    
+    def test_prelu(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        weight = torch.randn(1, device='cuda', dtype=self.dtype)
+        output = F.prelu(inp, weight)
+    
+    def test_rrelu(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.rrelu(inp, lower=1./8, upper=1./3, training=False, inplace=False)
+    
+    def test_rrelu_(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.rrelu(inp, lower=1./8, upper=1./3, training=False)
+    
+    def test_glu(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.glu(inp, dim=-1)
+    
+    def test_logsigmoid(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.logsigmoid(inp)
+    
+    def test_hardshrink(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.hardshrink(inp, lambd=0.5)
+    
+    def test_tanhshrink(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.tanhshrink(inp)
+    
+    def test_softsign(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.softsign(inp)
+    
+    def test_softplus(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.softplus(inp, beta=1, threshold=20)
+    
+    def test_softmin(self):
+        inp = torch.randn(16, 1024, device='cuda', dtype=self.dtype)
+        output = F.softmin(inp, dim=1,  _stacklevel=3, dtype=self.dtype)
+    
+    def test_softmax(self):
+        inp = torch.randn(16, 1024, device='cuda', dtype=self.dtype)
+        output = F.softmax(inp, dim=1, _stacklevel=3, dtype=self.dtype)
+    
+    def test_softshrink(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.softshrink(inp, lambd=0.5)
+    
+    def test_gumbel_softmax(self):
+        inp = torch.randn(16, 1024, device='cuda', dtype=self.dtype)
+        output = F.gumbel_softmax(inp, tau=1, hard=False, eps=1e-10, dim=-1)
+    
+    def test_log_softmax(self):
+        inp = torch.randn(16, 1024, device='cuda', dtype=self.dtype)
+        output = F.log_softmax(inp, dim=-1, _stacklevel=3)
+    
+    def test_tanh(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = torch.tanh(inp)
+    
+    def test_sigmoid(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = torch.sigmoid(inp)
+    
+    def test_batch_norm(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        # running_mean, running_var
+        running_mean = torch.randn(3, device='cuda', dtype=self.dtype)
+        running_var = torch.randn(3, device='cuda', dtype=self.dtype)
+        output = F.batch_norm(inp, running_mean, running_var, weight=None, bias=None, training=False, momentum=0.1, eps=1e-05)
+    
+    def test_instance_norm(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        running_mean = torch.randn(3, device='cuda', dtype=self.dtype)
+        running_var = torch.randn(3, device='cuda', dtype=self.dtype)
+        output = F.instance_norm(inp, running_mean=running_mean, running_var=running_var, weight=None, bias=None, use_input_stats=True, momentum=0.1, eps=1e-05)
+    
+    def test_layer_norm(self):
+        inp = torch.randn(1, 3, 32, 32, device='cuda', dtype=self.dtype)
+        output = F.layer_norm(inp, inp.size()[1:], weight=None, bias=None, eps=1e-05)
+    
+    def test_local_response_norm(self):
+        inp = torch.randn(16, 8, 64, 64, device='cuda', dtype=self.dtype)
+        output = F.local_response_norm(inp, 2, alpha=0.0001, beta=0.75, k=1.0)
+    
+    def test_normalize(self):
+        inp = torch.randn(16, 8, 64, 64, device='cuda', dtype=self.dtype)
+        output = F.normalize(inp, p=2, dim=1, eps=1e-12, out=None)
+    
+    def test_linear(self):
+        inp = torch.randn(32, 64, 128, device='cuda', dtype=self.dtype)
+        weight = torch.randn(256, 128, device='cuda', dtype=self.dtype)
+        output = F.linear(inp, weight, bias=None)
+    
+    def test_dropout(self):
+        inp = torch.randn(16, 8, 64, 64, device='cuda', dtype=self.dtype)
+        output = F.dropout(inp, p=0.5, training=True, inplace=False)
+    
+    def test_alpha_dropout(self):
+        inp = torch.randn(16, 8, 64, 64, device='cuda', dtype=self.dtype)
+        output = F.alpha_dropout(inp, p=0.5, training=True, inplace=False)
+    
+    def test_dropout2d(self):
+        inp = torch.randn(16, 8, 64, 64, device='cuda', dtype=self.dtype)
+        output = F.dropout2d(inp, p=0.5, training=True, inplace=False)
+    
+    def test_dropout3d(self):
+        inp = torch.randn(16, 8, 32, 64, 64, device='cuda', dtype=self.dtype)
+        output = F.dropout3d(inp, p=0.5, training=True, inplace=False)
+    
+    def test_embedding(self):
+        pre_embed_dim = 1024
+        post_embed_dim = 32
+        inp = torch.randint(0, pre_embed_dim, (128, 16), device='cuda')    
+        weight = torch.randn(pre_embed_dim, post_embed_dim, device='cuda', dtype=self.dtype)
+        output = F.embedding(inp, weight, padding_idx=None, max_norm=None, norm_type=2.0, scale_grad_by_freq=False, sparse=False)
+    
+    def test_embedding_bag(self):
+        pre_embed_dim = 1024
+        post_embed_dim = 32
+        inp = torch.randint(0, pre_embed_dim, (128, 16), device='cuda')    
+        weight = torch.randn(pre_embed_dim, post_embed_dim, device='cuda', dtype=self.dtype)
+        output = F.embedding_bag(inp, weight, offsets=None, max_norm=None, norm_type=2,
+            scale_grad_by_freq=False, mode='mean', sparse=False)
+    
+    def test_one_hot(self):
+        num_classes = 10
+        inp = torch.randint(0, num_classes, (128, 16), device='cuda')    
+        output = F.one_hot(inp, num_classes=10) 
+    
+    def test_pairwise_distance(self):
+        inp1 = torch.randn(1024, 128, device='cuda', dtype=self.dtype)
+        inp2 = torch.randn(1024, 128, device='cuda', dtype=self.dtype)
+        output = F.pairwise_distance(inp1, inp2, p=2.0, eps=1e-06, keepdim=False) 
+    
+    def test_cosine_similarity(self):
+        inp1 = torch.randn(1024, 128, device='cuda', dtype=self.dtype)
+        inp2 = torch.randn(1024, 128, device='cuda', dtype=self.dtype)
+        output = F.cosine_similarity(inp1, inp2, dim=1, eps=1e-8)
+    
+    def test_pdist(self):
+        # pdist is not implemented for fp16
+        inp = torch.randn(128, 128, device='cuda', dtype=torch.float32)
+        output = F.pdist(inp, p=2)
+    
+    def test_binary_cross_entropy(self):
+        # binary_cross_entropy is not implemented for fp16
+        inp = torch.randn(32, 128, device='cuda', dtype=torch.float32, requires_grad=True)
+        target = torch.randn(32, 128, device='cuda', dtype=torch.float32, requires_grad=False)
+        output = F.binary_cross_entropy(torch.sigmoid(inp), target)
+    
+    def test_binary_cross_entropy_with_logits(self):
+        inp = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = torch.empty_like(inp).random_(2)
+        output = F.binary_cross_entropy_with_logits(inp, target)
+    
+    def test_poisson_nll_loss(self):
+        inp = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=False)
+        output = F.poisson_nll_loss(inp, target, log_input=True, full=False,
+            size_average=None, eps=1e-08, reduce=None, reduction='mean')
+    
+    def test_cosine_embedding_loss(self):
+        inp1 = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        inp2 = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = torch.randn(32, device='cuda', dtype=self.dtype, requires_grad=False)
+        output = F.cosine_embedding_loss(inp1, inp2, target, margin=0,
+            size_average=None, reduce=None, reduction='mean')
+    
+    def test_cross_entropy(self):
+        inp = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = torch.randint(0, 100, (32,), device='cuda', dtype=torch.long, requires_grad=False)
+        output = F.cross_entropy(inp, target, weight=None, size_average=None,
+            ignore_index=-100, reduce=None, reduction='mean')
+    
+    def test_ctc_loss(self):
+        # force fp32 because _th_normal_ (used by next line is not supported for fp16)
+        log_probs = torch.randn(50, 16, 20, device='cuda', dtype=torch.float32).log_softmax(2).detach().requires_grad_()
+        targets = torch.randint(1, 20, (16, 30), device='cuda', dtype=torch.long)
+        input_lengths = torch.full((16,), 50, dtype=torch.long)
+        target_lengths = torch.randint(10, 30, (16,), dtype=torch.long)
+        loss = F.ctc_loss(log_probs, targets, input_lengths, target_lengths)
+    
+    def test_hinge_embedding_loss(self):
+        inp = torch.randn(128, 32, device='cuda', dtype=self.dtype)
+        target = torch.randint(0, 1, (32,), device='cuda') - 1
+        output = F.hinge_embedding_loss(inp, target, margin=1.0, size_average=None, reduce=None, reduction='mean') 
+    
+    def test_kl_div(self):
+        inp = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        output = F.kl_div(inp, target, size_average=None, reduce=None, reduction='batchmean')
+    
+    def test_mse_loss(self):
+        inp = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        output = F.mse_loss(inp, target, size_average=None, reduce=None, reduction='mean')
+    
+    def test_margin_ranking_loss(self):
+        inp1 = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        inp2 = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = (torch.randint(0, 1, (128,), device='cuda') - 1).type_as(inp1)
+        output = F.margin_ranking_loss(inp1, inp2, target, margin=0, size_average=None, reduce=None, reduction='mean')
+    
+    def test_multilabel_margin_loss(self):
+        inp = torch.randn(1024, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = torch.randint(0, 10, (1024,), dtype=torch.long, device='cuda')
+        output = F.multilabel_margin_loss(inp, target, size_average=None, reduce=None, reduction='mean')
+    
+    def test_nll_loss(self):
+        inp = torch.randn(64, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = torch.randint(0, 10, (64,), device='cuda', dtype=torch.long) 
+        output = F.nll_loss(inp, target, weight=None, size_average=None, ignore_index=-100, reduce=None, reduction='mean')
+    
+    def test_smooth_l1_loss(self):
+        inp = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=False)
+        output = F.smooth_l1_loss(inp, target, size_average=None, reduce=None, reduction='mean')
+    
+    def test_soft_margin_loss(self):
+        inp = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        target = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=False)
+        output = F.soft_margin_loss(inp, target, size_average=None, reduce=None, reduction='mean') 
+    
+    def test_triplet_margin_loss(self):
+        inp1 = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        inp2 = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        inp3 = torch.randn(32, 128, device='cuda', dtype=self.dtype, requires_grad=True)
+        output = F.triplet_margin_loss(inp1, inp2, inp3, margin=1.0, p=2,
+             eps=1e-06, swap=False, size_average=None, reduce=None, reduction='mean')
+    
+    def test_pixel_shuffle(self):
+        inp = torch.randn(16, 8, 64, 64, device='cuda', dtype=self.dtype)
+        output = torch.nn.functional.pixel_shuffle(inp, 2)
+    
+    def test_pad(self):
+        inp = torch.randn(16, 8, 64, 64, device='cuda', dtype=self.dtype)
+        pad = (3, 3)
+        output = F.pad(inp, pad, mode='constant', value=0)
+    
+    def test_interpolate(self):
+        inp = torch.randn(16, 8, 64, 64, device='cuda', dtype=self.dtype)
+        output = F.interpolate(inp, size=None, scale_factor=2, mode='nearest', align_corners=None)
+    
+    def test_grid_sample(self):
+        inp = torch.randn(16, 8, 64, 64, device='cuda', dtype=self.dtype)
+        grid = torch.randn(16, 32, 32, 2, device='cuda', dtype=self.dtype)
+        output = F.grid_sample(inp, grid, mode='bilinear', padding_mode='zeros')
+    
+    def test_affine_grid(self):
+        theta = torch.randn(32, 2, 3, device='cuda', dtype=self.dtype)
+        size = (32, 8, 32, 32)
+        output = F.affine_grid(theta, size)
+
+
+def run_tests(precision):
+    dummy = TestPyProfNvtx('test_affine_grid', None)
+    test_cases = list(filter(lambda x: 'test_' in x, map(lambda x: x[0], inspect.getmembers(dummy, predicate=inspect.ismethod))))
+    print("Running tests for {}".format(precision))
+    suite = unittest.TestSuite()
+    for test_case in test_cases:
+        suite.addTest(TestPyProfNvtx(test_case, precision))
+    unittest.TextTestRunner().run(suite)
+
+if __name__ == '__main__':
+    run_tests(torch.float32)
+    run_tests(torch.float16)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_test.py
new file mode 100644
index 0000000000..14527a8919
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_test.py
@@ -0,0 +1,72 @@
+"""L0 Tests Runner.
+
+How to run this script?
+
+1. Run all the tests: `python /path/to/apex/tests/L0/run_test.py`
+2. Run one of the tests (e.g. fused layer norm):
+    `python /path/to/apex/tests/L0/run_test.py --include run_fused_layer_norm`
+3. Run two or more of the tests (e.g. optimizers and fused layer norm):
+    `python /path/to/apex/tests/L0/run_test.py --include run_optimizers run_fused_layer_norm`
+"""
+import argparse
+import os
+import unittest
+import sys
+
+
+TEST_ROOT = os.path.dirname(os.path.abspath(__file__))
+TEST_DIRS = [
+    "run_amp",
+    "run_fp16util",
+    "run_optimizers",
+    "run_fused_layer_norm",
+    "run_pyprof_nvtx",
+    "run_pyprof_data",
+    "run_mlp",
+    "run_transformer",
+]
+DEFAULT_TEST_DIRS = [
+    "run_optimizers",
+    "run_fused_layer_norm",
+    "run_mlp",
+    "run_transformer",
+]
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(
+        description="L0 test runner",
+        formatter_class=argparse.ArgumentDefaultsHelpFormatter,
+    )
+    parser.add_argument(
+        "--include",
+        nargs="+",
+        choices=TEST_DIRS,
+        default=DEFAULT_TEST_DIRS,
+        help="select a set of tests to run (defaults to ALL tests).",
+    )
+    args, _ = parser.parse_known_args()
+    return args
+
+
+def main(args):
+    runner = unittest.TextTestRunner(verbosity=2)
+    errcode = 0
+    for test_dir in args.include:
+        test_dir = os.path.join(TEST_ROOT, test_dir)
+        print(test_dir)
+        suite = unittest.TestLoader().discover(test_dir)
+
+        print("\nExecuting tests from " + test_dir)
+
+        result = runner.run(suite)
+
+        if not result.wasSuccessful():
+            errcode = 1
+
+    sys.exit(errcode)
+
+
+if __name__ == '__main__':
+    args = parse_args()
+    main(args)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_cross_entropy_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_cross_entropy_test.py
new file mode 100644
index 0000000000..0a0f784ca3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_cross_entropy_test.py
@@ -0,0 +1,110 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import torch
+import torch.nn.functional as F
+
+from apex.transformer import parallel_state
+from apex.transformer import tensor_parallel
+from apex.transformer.tensor_parallel.cross_entropy import vocab_parallel_cross_entropy
+from apex.transformer.testing import global_vars
+from apex.transformer.testing.commons import set_random_seed
+from apex.transformer.testing.commons import IdentityLayer
+from apex.transformer.testing.commons import print_separator
+from apex.transformer.testing.commons import initialize_distributed
+from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE
+
+
+global_vars.set_global_variables()
+
+
+def torch_cross_entropy(batch_size, seq_length, vocab_size,
+                        logits_scale, seed):
+    set_random_seed(seed)
+    identity = IdentityLayer((batch_size, seq_length, vocab_size),
+                             scale=logits_scale).cuda()
+    logits = identity()
+    target = torch.cuda.LongTensor(
+        size=(batch_size, seq_length)).random_(0, vocab_size)
+    loss = F.cross_entropy(logits.view(-1, logits.size()[-1]),
+                           target.view(-1),
+                           reduction='none').view_as(target).mean()
+    loss.backward()
+    return loss, identity.weight.grad
+
+
+def tensor_sharded_cross_entropy(batch_size, seq_length, vocab_size, logits_scale, seed):
+    set_random_seed(seed)
+    identity = IdentityLayer((batch_size, seq_length, vocab_size), scale=logits_scale).cuda()
+    logits = identity()
+    logits_parallel = tensor_parallel.scatter_to_tensor_model_parallel_region(logits)
+    target = torch.cuda.LongTensor(
+        size=(batch_size, seq_length)).random_(0, vocab_size)
+    logits_parallel_ = logits_parallel.clone().detach()
+    loss = vocab_parallel_cross_entropy(logits_parallel, target).mean()
+    loss.backward()
+    # check for mutation
+    assert torch.equal(logits_parallel_, logits_parallel)
+    return loss, identity.weight.grad
+
+
+def test_cross_entropy(tensor_model_parallel_size):
+
+    if torch.distributed.get_rank() == 0:
+        print('> testing cross entropy with model parallel size {} ...'.
+              format(tensor_model_parallel_size))
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    batch_size = 13
+    seq_length = 17
+    vocab_size_per_partition = 11
+    logits_scale = 1000.0
+    vocab_size = vocab_size_per_partition * tensor_model_parallel_size
+    seed = 1234
+
+    loss_torch, grad_torch = torch_cross_entropy(batch_size, seq_length, vocab_size, logits_scale, seed)
+    loss_mpu, grad_mpu = tensor_sharded_cross_entropy(batch_size, seq_length, vocab_size, logits_scale, seed)
+
+    error = loss_torch.sub_(loss_mpu).abs().max()
+    print('   max error in loss on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    error = grad_torch.sub_(grad_mpu).abs().max()
+    print('   max error in grad on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(TEST_SUCCESS_MESSAGE)
+
+
+if __name__ == '__main__':
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+    initialize_distributed()
+    world_size = torch.distributed.get_world_size()
+
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        print_separator('test cross entropy')
+        test_cross_entropy(tensor_model_parallel_size)
+        tensor_model_parallel_size *= 2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_data_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_data_test.py
new file mode 100644
index 0000000000..981efeba93
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_data_test.py
@@ -0,0 +1,95 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import functools
+import operator
+
+import torch
+
+from apex.transformer import parallel_state
+from apex.transformer.tensor_parallel import data as data_utils
+from apex.transformer.testing import global_vars
+from apex.transformer.testing.commons import print_separator
+from apex.transformer.testing.commons import initialize_distributed
+from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE
+
+global_vars.set_global_variables()
+
+
+def test_broadcast_data(tensor_model_parallel_size):
+
+    if torch.distributed.get_rank() == 0:
+        print('> testing broadcast_data with model parallel size {} ...'.
+              format(tensor_model_parallel_size))
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    torch.manual_seed(1234 + parallel_state.get_data_parallel_rank())
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    key_size_t = {
+        'key1': [7, 11],
+        'key2': [8, 2, 1],
+        'key3': [13],
+        'key4': [5, 1, 2],
+        'key5': [5, 12],
+    }
+    keys = list(key_size_t.keys())
+
+    data = {}
+    data_t = {}
+    for key in key_size_t:
+        data[key] = torch.LongTensor(size=key_size_t[key]).random_(0, 1000)
+        data_t[key] = data[key].clone()
+    data['keyX'] = torch.FloatTensor(size=(5, )).random_(0, 1000)
+    data_t['keyX'] = data['keyX'].clone()
+    if parallel_state.get_tensor_model_parallel_rank() != 0:
+        data = None
+
+    data_utils._check_data_types(keys, data_t, torch.int64)
+    key_size, key_numel, \
+        total_numel = data_utils._build_key_size_numel_dictionaries(keys, data)
+    for key in keys:
+        assert key_size[key] == key_size_t[key]
+    total_numel_t = 0
+    for key in keys:
+        target_size = functools.reduce(operator.mul, key_size_t[key], 1)
+        assert key_numel[key] == target_size
+        total_numel_t += target_size
+    assert total_numel == total_numel_t
+
+    data_b = data_utils.broadcast_data(keys, data, torch.int64)
+    for key in keys:
+        tensor = data_t[key].cuda()
+        assert data_b[key].sub(tensor).abs().max() == 0
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(TEST_SUCCESS_MESSAGE)
+
+
+if __name__ == '__main__':
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+    initialize_distributed()
+    world_size = torch.distributed.get_world_size()
+
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        print_separator('test test broadcast data')
+        test_broadcast_data(tensor_model_parallel_size)
+        tensor_model_parallel_size *= 2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_initialize_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_initialize_test.py
new file mode 100644
index 0000000000..492fa04a35
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_initialize_test.py
@@ -0,0 +1,104 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import torch
+
+from apex.transformer import parallel_state
+from apex.transformer.testing import global_vars
+from apex.transformer.testing.commons import print_separator
+from apex.transformer.testing.commons import initialize_distributed
+from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE
+
+
+global_vars.set_global_variables()
+
+
+def test_initialize_model_parallel(tensor_model_parallel_size):
+
+    if torch.distributed.get_rank() == 0:
+        print('> testing initialize_model_parallel with size {} ...'.format(
+            tensor_model_parallel_size))
+    tensor_model_parallel_size_ = min(
+        tensor_model_parallel_size,
+        torch.distributed.get_world_size(),
+    )
+    assert not parallel_state.model_parallel_is_initialized()
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size_)
+    assert parallel_state.model_parallel_is_initialized()
+
+    # Checks.
+    def check(group, world_size, rank):
+        assert world_size == torch.distributed.get_world_size(group=group)
+        assert rank == torch.distributed.get_rank(group=group)
+
+    # Model parallel.
+    world_size = tensor_model_parallel_size_
+    rank = torch.distributed.get_rank() % tensor_model_parallel_size_
+    assert world_size == parallel_state.get_tensor_model_parallel_world_size()
+    assert rank == parallel_state.get_tensor_model_parallel_rank()
+    check(parallel_state.get_tensor_model_parallel_group(), world_size, rank)
+
+    # Data parallel.
+    world_size = torch.distributed.get_world_size() // tensor_model_parallel_size_
+    rank = torch.distributed.get_rank() // tensor_model_parallel_size
+    assert world_size == parallel_state.get_data_parallel_world_size()
+    assert rank == parallel_state.get_data_parallel_rank()
+    check(parallel_state.get_data_parallel_group(), world_size, rank)
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(TEST_SUCCESS_MESSAGE)
+
+
+def test_get_tensor_model_parallel_src_rank(tensor_model_parallel_size_):
+
+    if torch.distributed.get_rank() == 0:
+        print('> testing get_tensor_model_parallel_src_rank with size {} ...'.format(
+            tensor_model_parallel_size_))
+    tensor_model_parallel_size = min(
+        tensor_model_parallel_size_,
+        torch.distributed.get_world_size(),
+    )
+    assert not parallel_state.model_parallel_is_initialized()
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    assert parallel_state.model_parallel_is_initialized()
+
+    # Checks
+    src_rank = torch.distributed.get_rank() - parallel_state.get_tensor_model_parallel_rank()
+    assert parallel_state.get_tensor_model_parallel_src_rank() == src_rank
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print('>> passed the test :-)')
+
+
+if __name__ == '__main__':
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+    initialize_distributed()
+    world_size = torch.distributed.get_world_size()
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        print_separator('test initialize model parallel')
+        test_initialize_model_parallel(tensor_model_parallel_size)
+        print_separator('test model parallel source rank')
+        test_get_tensor_model_parallel_src_rank(tensor_model_parallel_size)
+        tensor_model_parallel_size *= 2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_layers_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_layers_test.py
new file mode 100644
index 0000000000..25ba48282c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_layers_test.py
@@ -0,0 +1,696 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import torch
+import torch.nn.init as init
+from torch.nn.parameter import Parameter
+
+from apex.transformer import parallel_state
+from apex.transformer.tensor_parallel import layers
+from apex.transformer.testing import global_vars
+from apex.transformer.testing.commons import set_random_seed
+from apex.transformer.testing.commons import print_separator
+from apex.transformer.testing.commons import initialize_distributed
+from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE
+
+
+global_vars.set_global_variables()
+
+
+class IdentityLayer3D(torch.nn.Module):
+    def __init__(self, m, n, k):
+        super(IdentityLayer3D, self).__init__()
+        self.weight = Parameter(torch.Tensor(m, n, k))
+        torch.nn.init.xavier_normal_(self.weight)
+
+    def forward(self):
+        return self.weight
+
+
+def test_parallel_embedding(tensor_model_parallel_size):
+
+    if torch.distributed.get_rank() == 0:
+        print('> testing parallel embedding with model parallel size {} ...'.
+              format(tensor_model_parallel_size))
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    batch_size = 17
+    seq_length = 23
+    vocab_size = 48
+    hidden_size = 16
+    seed = 1236
+
+    set_random_seed(123)
+    input_data = torch.LongTensor(
+        size=(batch_size, seq_length)).random_(0, vocab_size).cuda()
+    loss_weight = torch.randn([batch_size, seq_length, hidden_size]).cuda()
+
+    set_random_seed(seed)
+    embedding_original = torch.nn.Embedding(vocab_size, hidden_size).cuda()
+
+    output = embedding_original(input_data)
+    loss_original = torch.mul(output, loss_weight).sum()
+    loss_original.backward()
+
+    set_random_seed(seed)
+    embedding_parallel = layers.ParallelEmbedding(
+        vocab_size, hidden_size, init_method=init.normal_).cuda()
+    output = embedding_parallel(input_data)
+    loss_parallel = torch.mul(output, loss_weight).sum()
+    loss_parallel.backward()
+
+    set_random_seed(seed)
+    embedding_vocab_parallel = layers.VocabParallelEmbedding(
+        vocab_size, hidden_size, init_method=init.normal_).cuda()
+    output = embedding_vocab_parallel(input_data)
+    loss_vocab_parallel = torch.mul(output, loss_weight).sum()
+    loss_vocab_parallel.backward()
+
+    torch.distributed.barrier()
+    error = loss_parallel.sub(loss_original).abs()
+    print('   error in loss (parallel) on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-12, 'error: {}'.format(error)
+
+    torch.distributed.barrier()
+    error = loss_vocab_parallel.sub(loss_original).abs()
+    print('   error in loss (vocab parallel) on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-12, 'error: {}'.format(error)
+
+    weight_grad_orig = torch.split(embedding_original.weight.grad,
+                                   hidden_size // tensor_model_parallel_size,
+                                   1)[parallel_state.get_tensor_model_parallel_rank()]
+    error = embedding_parallel.weight.grad.sub(weight_grad_orig).abs().max()
+    print('   error in grad (parallel) on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-12, 'error: {}'.format(error)
+
+    weight_grad_orig = torch.split(embedding_original.weight.grad,
+                                   vocab_size // tensor_model_parallel_size,
+                                   0)[parallel_state.get_tensor_model_parallel_rank()]
+    error = embedding_vocab_parallel.weight.grad.sub(
+        weight_grad_orig).abs().max()
+    print('   error in grad (vocab parallel) on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-12, 'error: {}'.format(error)
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print('>> passed the test :-)')
+
+
+def test_initialize_affine_weight(tensor_model_parallel_size, device):
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    if torch.distributed.get_rank() == 0:
+        print('> testing initialize_affine_weight with model parallel '
+              'size: {}'.format(tensor_model_parallel_size))
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    seed = 12345
+    input_size_coeff = 13
+    input_size = input_size_coeff * tensor_model_parallel_size
+    output_size_coeff = 17
+    output_size = output_size_coeff * tensor_model_parallel_size
+
+    # ---------------
+    # Column parallel
+    # ---------------
+    weight = torch.empty(output_size_coeff, input_size)
+    set_random_seed(seed)
+    if device == 'cpu':
+        layers._initialize_affine_weight_cpu(weight, output_size, input_size,
+                                             output_size_coeff, 0,
+                                             torch.nn.init.normal_,
+                                             params_dtype=global_vars.get_args().params_dtype,
+                                             )
+    else:
+        layers._initialize_affine_weight_gpu(weight, torch.nn.init.normal_, 0)
+
+    # Target.
+    set_random_seed(seed)
+    master_weight = torch.empty(output_size, input_size)
+    torch.nn.init.normal_(master_weight)
+    rank = parallel_state.get_tensor_model_parallel_rank()
+    my_weight = torch.split(master_weight, output_size_coeff,
+                            dim=0)[rank].contiguous().clone()
+
+    # Compare.
+    error = weight.sub(my_weight).abs().max()
+    torch.distributed.barrier()
+    print('   column parallel max error (should be zero) on global rank '
+          '{}: {}'.format(torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    # ------------
+    # Row parallel
+    # ------------
+    weight = torch.empty(output_size, input_size_coeff)
+    set_random_seed(seed)
+    if device == 'cpu':
+        layers._initialize_affine_weight_cpu(
+            weight, output_size, input_size, input_size_coeff, 1, torch.nn.init.normal_,
+            params_dtype=global_vars.get_args().params_dtype)
+
+    else:
+        layers._initialize_affine_weight_gpu(weight, torch.nn.init.normal_, 1)
+
+    # Target.
+    set_random_seed(seed)
+    master_weight = torch.empty(output_size, input_size)
+    torch.nn.init.normal_(master_weight)
+    rank = parallel_state.get_tensor_model_parallel_rank()
+    my_weight = torch.split(master_weight, input_size_coeff,
+                            dim=1)[rank].contiguous().clone()
+
+    # Compare.
+    error = weight.sub(my_weight).abs().max()
+    torch.distributed.barrier()
+    print('   row parallel max error (should be zero) on global rank '
+          '{}: {}'.format(torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(' >> passed the test :-)')
+
+
+class IdentityLayer2D(torch.nn.Module):
+    def __init__(self, m, n):
+        super(IdentityLayer2D, self).__init__()
+        self.weight = Parameter(torch.Tensor(m, n))
+        torch.nn.init.xavier_normal_(self.weight)
+
+    def forward(self):
+        return self.weight
+
+
+def test_column_parallel_linear(tensor_model_parallel_size):
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    if torch.distributed.get_rank() == 0:
+        print('> testing ColumnParallelLinear with model parallel '
+              'size: {}'.format(tensor_model_parallel_size))
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    seed = 12345
+    set_random_seed(seed)
+    input_size_coeff = 13
+    input_size = input_size_coeff * tensor_model_parallel_size
+    output_size_coeff = 17
+    output_size = output_size_coeff * tensor_model_parallel_size
+    batch_size = 7
+
+    # Network
+    identity_layer = IdentityLayer2D(batch_size, input_size).cuda()
+    linear_layer = layers.ColumnParallelLinear(
+        input_size, output_size, keep_master_weight_for_test=True,
+        params_dtype=global_vars.get_args().params_dtype,
+        use_cpu_initialization=global_vars.get_args().use_cpu_initialization,
+    ).cuda()
+    loss_weight = torch.randn([batch_size, output_size]).cuda()
+    # Forward
+    input_ = identity_layer()
+    output, _ = linear_layer(input_)
+    loss = torch.mul(output, loss_weight).sum()
+    # Backward
+    loss.backward()
+
+    # Values.
+    dLdY = loss_weight
+    X = identity_layer.weight
+    A = linear_layer.master_weight.cuda()
+    dLdA = torch.matmul(dLdY.t(), X)
+    dLdb = torch.matmul(torch.ones(batch_size, 1).cuda().t(), dLdY).view(-1)
+    dLdX = torch.matmul(dLdY, A)
+
+    rank = parallel_state.get_tensor_model_parallel_rank()
+    my_dLdA = torch.split(dLdA, output_size_coeff,
+                          dim=0)[rank].contiguous().clone()
+    error = my_dLdA.sub(linear_layer.weight.grad).abs().max()
+    torch.distributed.barrier()
+    print('   error in dLdA on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    my_dLdb = torch.split(dLdb, output_size_coeff,
+                          dim=0)[rank].contiguous().clone()
+    error = my_dLdb.sub(linear_layer.bias.grad).abs().max()
+    torch.distributed.barrier()
+    print('   error in dLdb on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    error = dLdX.sub(identity_layer.weight.grad).abs().max()
+    torch.distributed.barrier()
+    print('   error in dLdX on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(' >> passed the test :-)')
+
+
+def test_column_parallel_linear_with_async_allreduce_autocast(tensor_model_parallel_size):
+    autocast_dtypes = (torch.half, torch.bfloat16) if torch.cuda.is_bf16_supported() else (torch.half,)
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    seed = 12345
+    set_random_seed(seed)
+    input_size_coeff = 13
+    input_size = input_size_coeff * tensor_model_parallel_size
+    output_size_coeff = 17
+    output_size = output_size_coeff * tensor_model_parallel_size
+    batch_size = 7
+
+    # Network
+    identity_layer = IdentityLayer3D(batch_size, batch_size, input_size).cuda()
+    linear_layer = layers.ColumnParallelLinear(
+        input_size, output_size, keep_master_weight_for_test=True,
+        params_dtype=global_vars.get_args().params_dtype,
+        use_cpu_initialization=global_vars.get_args().use_cpu_initialization,
+    ).cuda()
+    assert linear_layer.async_tensor_model_parallel_allreduce or tensor_model_parallel_size == 1
+    # Forward
+    for dtype in autocast_dtypes:
+        loss_weight = torch.randn([batch_size, output_size]).cuda()
+        with torch.cuda.amp.autocast(dtype=dtype):
+            output, _ = linear_layer(identity_layer())
+            loss = torch.mul(output, loss_weight).sum()
+        assert output.dtype == dtype
+        # Backward
+        loss.backward()
+        torch.distributed.barrier()
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(' >> passed the test :-)')
+
+
+def test_column_parallel_linear_with_async_allreduce_custom_amp(tensor_model_parallel_size):
+    dtypes = (torch.half, torch.bfloat16) if torch.cuda.is_bf16_supported() else (torch.half,)
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    seed = 12345
+    set_random_seed(seed)
+    input_size_coeff = 13
+    input_size = input_size_coeff * tensor_model_parallel_size
+    output_size_coeff = 17
+    output_size = output_size_coeff * tensor_model_parallel_size
+    batch_size = 7
+
+    for dtype in dtypes:
+        # Network
+        identity_layer = IdentityLayer3D(batch_size, batch_size, input_size).to(device="cuda", dtype=dtype)
+        linear_layer = layers.ColumnParallelLinear(
+            input_size, output_size, keep_master_weight_for_test=True,
+            params_dtype=global_vars.get_args().params_dtype,
+            use_cpu_initialization=global_vars.get_args().use_cpu_initialization,
+        ).to(device="cuda", dtype=dtype)
+        # Forward
+        loss_weight = torch.randn([batch_size, output_size]).cuda()
+        output, _ = linear_layer(identity_layer())
+        loss = torch.mul(output, loss_weight).sum()
+        loss.backward()
+        torch.distributed.barrier()
+
+        assert output.dtype == dtype
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(' >> passed the test :-)')
+
+
+def test_row_parallel_linear(tensor_model_parallel_size):
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    if torch.distributed.get_rank() == 0:
+        print('> testing RowParallelLinear with model parallel '
+              'size: {}'.format(tensor_model_parallel_size))
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    seed = 12345
+    set_random_seed(seed)
+    input_size_coeff = 13
+    input_size = input_size_coeff * tensor_model_parallel_size
+    output_size_coeff = 17
+    output_size = output_size_coeff * tensor_model_parallel_size
+    batch_size = 7
+
+    # Network
+    identity_layer = IdentityLayer2D(batch_size, input_size).cuda()
+    linear_layer = layers.RowParallelLinear(
+        input_size, output_size, keep_master_weight_for_test=True,
+        params_dtype=global_vars.get_args().params_dtype,
+        use_cpu_initialization=global_vars.get_args().use_cpu_initialization,
+    ).cuda()
+    loss_weight = torch.randn([batch_size, output_size]).cuda()
+    # Forward
+    input_ = identity_layer()
+    output, _ = linear_layer(input_)
+    loss = torch.mul(output, loss_weight).sum()
+    # Backward
+    loss.backward()
+
+    # Values.
+    dLdY = loss_weight
+    X = identity_layer.weight
+    A = linear_layer.master_weight.cuda()
+    dLdA = torch.matmul(dLdY.t(), X)
+    dLdb = torch.matmul(torch.ones(batch_size, 1).cuda().t(), dLdY).view(-1)
+    dLdX = torch.matmul(dLdY, A)
+
+    rank = parallel_state.get_tensor_model_parallel_rank()
+    my_dLdA = torch.split(dLdA, input_size_coeff,
+                          dim=1)[rank].contiguous().clone()
+    error = my_dLdA.sub(linear_layer.weight.grad).abs().max()
+    torch.distributed.barrier()
+    print('   error in dLdA on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    error = dLdb.sub(linear_layer.bias.grad).abs().max()
+    torch.distributed.barrier()
+    print('   error in dLdb on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    error = dLdX.sub(identity_layer.weight.grad).abs().max()
+    torch.distributed.barrier()
+    print('   error in dLdX on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(' >> passed the test :-)')
+
+
+def parallel_self_attention(tensor_model_parallel_size, num_att_heads_per_partition,
+                            hidden_size_per_att_head, dropout_prob, batch_size,
+                            sequence_length):
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    seed = 12345
+    set_random_seed(seed)
+
+    num_att_heads = num_att_heads_per_partition * \
+        torch.distributed.get_world_size()
+    hidden_size = hidden_size_per_att_head * num_att_heads
+
+    # Network
+    identity_layer = IdentityLayer3D(batch_size, sequence_length,
+                                     hidden_size).cuda()
+    attention_layer = parallel_state.BertParallelSelfAttention(hidden_size, num_att_heads,
+                                                    dropout_prob).cuda()
+    loss_weight = torch.randn([batch_size, sequence_length, hidden_size]).cuda()
+    attention_mask = torch.randn([batch_size, 1, 1, sequence_length]).cuda()
+    # Forward
+    input_ = identity_layer()
+    output = attention_layer(input_, attention_mask)
+    loss = torch.mul(output, loss_weight).sum()
+    # Backward
+    loss.backward()
+
+    rank = parallel_state.get_tensor_model_parallel_rank()
+    parallel_state.destroy_model_parallel()
+    return rank, hidden_size, tensor_model_parallel_size, loss, \
+        attention_layer, identity_layer
+
+
+def test_parallel_self_attention(tensor_model_parallel_size):
+
+    if torch.distributed.get_rank() == 0:
+        print('> testing ParallelSelfAttention with model parallel '
+              'size: {}'.format(tensor_model_parallel_size))
+
+    num_att_heads_per_partition = 3
+    hidden_size_per_att_head = 7
+    dropout_prob = 0.0  # has to be zero
+    batch_size = 5
+    sequence_length = 13
+
+    rank_1, hideen_size_1, tensor_model_parallel_size_1, loss_1, \
+        attention_layer_1, identity_layer_1 = parallel_self_attention(
+            1, num_att_heads_per_partition,
+            hidden_size_per_att_head, dropout_prob, batch_size, sequence_length)
+
+    rank, hidden_size, tensor_model_parallel_size, loss, \
+        attention_layer, identity_layer = parallel_self_attention(
+            tensor_model_parallel_size, num_att_heads_per_partition,
+            hidden_size_per_att_head, dropout_prob, batch_size, sequence_length)
+    assert hideen_size_1 == hidden_size
+
+    error = loss_1.sub(loss).abs().max()
+    torch.distributed.barrier()
+    print('   loss error on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 5.0e-6
+
+    my_lin_grad_list = torch.split(
+        attention_layer_1.query_key_value.weight.grad,
+        hidden_size // tensor_model_parallel_size, 0)[rank::tensor_model_parallel_size]
+    my_lin_grad = torch.cat(my_lin_grad_list, dim=0)
+    error = my_lin_grad.sub(
+        attention_layer.query_key_value.weight.grad).abs().max()
+    torch.distributed.barrier()
+    print('   weight gradient error on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 5.0e-6
+
+    error = identity_layer_1.weight.grad.sub(
+        identity_layer.weight.grad).abs().max()
+    torch.distributed.barrier()
+    print('   input gradient error on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 5.0e-6
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(' >> passed the test :-)')
+
+
+def parallel_transformer(tensor_model_parallel_size, num_att_heads_per_partition,
+                         hidden_size_per_att_head, batch_size, sequence_length):
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    seed = 12345
+    set_random_seed(seed)
+
+    num_att_heads = num_att_heads_per_partition * \
+        torch.distributed.get_world_size()
+    hidden_size = hidden_size_per_att_head * num_att_heads
+    intermediate_size = 4 * hidden_size
+
+    # Network
+    identity_layer = IdentityLayer3D(batch_size, sequence_length,
+                                     hidden_size).cuda()
+    transformer_layer = parallel_state.BertParallelTransformerLayer(
+        hidden_size, intermediate_size, num_att_heads, 0.0, 0.0,
+        torch.nn.functional.relu, 1.0e-5).cuda()
+
+    loss_weight = torch.randn([batch_size, sequence_length, hidden_size]).cuda()
+    attention_mask = torch.randn([batch_size, 1, 1, sequence_length]).cuda()
+    # Forward
+    input_ = identity_layer()
+    output = transformer_layer(input_, attention_mask)
+    loss = torch.mul(output, loss_weight).sum()
+    # Backward
+    loss.backward()
+
+    rank = parallel_state.get_tensor_model_parallel_rank()
+    parallel_state.destroy_model_parallel()
+    return rank, hidden_size, tensor_model_parallel_size, loss, \
+        transformer_layer, identity_layer
+
+
+def test_parallel_transformer_layer(tensor_model_parallel_size):
+
+    if torch.distributed.get_rank() == 0:
+        print('> testing ParallelTransformerLayer with model parallel '
+              'size: {}'.format(tensor_model_parallel_size))
+
+    num_att_heads_per_partition = 3
+    hidden_size_per_att_head = 7
+    batch_size = 5
+    sequence_length = 13
+
+    rank_1, hidden_size_1, tensor_model_parallel_size_1, loss_1, \
+        transformer_layer_1, identity_layer_1 = parallel_transformer(
+            1, num_att_heads_per_partition,
+            hidden_size_per_att_head, batch_size, sequence_length)
+
+    rank, hidden_size, tensor_model_parallel_size, loss, \
+        transformer_layer, identity_layer = parallel_transformer(
+            tensor_model_parallel_size, num_att_heads_per_partition,
+            hidden_size_per_att_head, batch_size, sequence_length)
+
+    error = loss_1.sub(loss).abs().max()
+    torch.distributed.barrier()
+    print('   loss error on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 5.0e-5, 'error: {}'.format(error)
+
+    error = identity_layer_1.weight.grad.sub(
+        identity_layer.weight.grad).abs().max()
+    torch.distributed.barrier()
+    print('   input gradient error on global rank {}: {}'.format(
+        torch.distributed.get_rank(), error))
+    assert error < 5.0e-5, 'error: {}'.format(error)
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(TEST_SUCCESS_MESSAGE)
+
+
+if __name__ == '__main__':
+    torch.backends.cudnn.deterministic = True
+    torch.backends.cudnn.benchmark = False
+
+    initialize_distributed()
+    world_size = torch.distributed.get_world_size()
+
+    exceptions = []
+
+    print_separator('test initialize affine weight cpu')
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        try:
+            test_initialize_affine_weight(tensor_model_parallel_size, 'cpu')
+        except Exception as e:
+            exceptions.append(f"test_initialize_affine_weight-cpu with tensor model parallel size of {tensor_model_parallel_size} failed: {str(e)}")
+            # Reset groups
+            parallel_state.destroy_model_parallel()
+            break
+        else:
+            tensor_model_parallel_size *= 2
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    print_separator('test initialize affine weight gpu')
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        try:
+            test_initialize_affine_weight(tensor_model_parallel_size, 'gpu')
+        except Exception as e:
+            exceptions.append(f"test_initialize_affine_weight-gpu with tensor model parallel size of {tensor_model_parallel_size} failed: {str(e)}")
+            # Reset groups
+            parallel_state.destroy_model_parallel()
+            break
+        else:
+            tensor_model_parallel_size *= 2
+
+    # Deleted, replaced with vocab parallel embedding?
+    #tensor_model_parallel_size = 1
+    #while tensor_model_parallel_size <= world_size:
+    #    print_separator('test parallel embedding')
+    #    test_parallel_embedding(tensor_model_parallel_size)
+    #    tensor_model_parallel_size *= 2
+
+    print_separator('test column-parallel linear')
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        try:
+            test_column_parallel_linear(tensor_model_parallel_size)
+        except Exception as e:
+            exceptions.append(f"test_column_parallel_linear with tensor model parallel size of {tensor_model_parallel_size} failed: {str(e)}")
+            # Reset groups
+            parallel_state.destroy_model_parallel()
+            break
+        else:
+            tensor_model_parallel_size *= 2
+
+    print_separator('test row-parallel linear')
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        try:
+            test_row_parallel_linear(tensor_model_parallel_size)
+        except Exception as e:
+            exceptions.append(f"test_row_parallel_linear with tensor model parallel size of {tensor_model_parallel_size} failed: {str(e)}")
+            # Reset groups
+            parallel_state.destroy_model_parallel()
+            break
+        else:
+            tensor_model_parallel_size *= 2
+
+    print_separator("test ColumnParallelLinearWithAsyncAllreduce - autocast")
+    tensor_model_parallel_size = 2
+    while tensor_model_parallel_size <= world_size:
+        try:
+            test_column_parallel_linear_with_async_allreduce_autocast(tensor_model_parallel_size)
+        except Exception as e:
+            exceptions.append(f"test_column_parallel_linear_with_async_allreduce_autocast with tensor model parallel size of {tensor_model_parallel_size} failed: {str(e)}")
+            # Reset groups
+            parallel_state.destroy_model_parallel()
+            break
+        else:
+            tensor_model_parallel_size *= 2
+
+    print_separator("test ColumnParallelLinearWithAsyncAllreduce - custom AMP")
+    tensor_model_parallel_size = 2
+    while tensor_model_parallel_size <= world_size:
+        try:
+            test_column_parallel_linear_with_async_allreduce_custom_amp(tensor_model_parallel_size)
+        except Exception as e:
+            exceptions.append(f"test_column_parallel_linear_with_async_allreduce_custom_amp with tensor model parallel size of {tensor_model_parallel_size} failed: {str(e)}")
+            # Reset groups
+            parallel_state.destroy_model_parallel()
+            break
+        else:
+            tensor_model_parallel_size *= 2
+
+    if exceptions:
+        raise RuntimeError("\n".join(exceptions))
+    # Deleted
+    #print_separator('test parallel self-attention')
+    #tensor_model_parallel_size = 1
+    #while tensor_model_parallel_size <= world_size:
+    #    test_parallel_self_attention(tensor_model_parallel_size)
+    #    tensor_model_parallel_size *= 2
+
+    #Deleted because PararallelTransformerLayer no longer exists
+    # print_separator('test parallel transformer')
+    # tensor_model_parallel_size = 1
+    # while tensor_model_parallel_size <= world_size:
+    #     test_parallel_transformer_layer(tensor_model_parallel_size)
+    #     tensor_model_parallel_size *= 2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_mappings_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_mappings_test.py
new file mode 100644
index 0000000000..b75118d67a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_mappings_test.py
@@ -0,0 +1,63 @@
+import torch
+
+from apex.transformer import parallel_state
+from apex.transformer.tensor_parallel import mappings
+from apex.transformer.testing import global_vars
+from apex.transformer.testing.commons import initialize_distributed
+
+global_vars.set_global_variables()
+
+
+def test__reduce(args, tensor_model_parallel_size):
+    print("Testing reduction size =", tensor_model_parallel_size)
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+    assert torch.equal(
+        mappings._reduce(torch.full((10, 10, 10, 10), (50))),
+        torch.full((10, 10, 10, 10), 50 * tensor_model_parallel_size),
+    )
+    parallel_state.destroy_model_parallel()
+    print("Passed!")
+
+
+def test__split(args, tensor_model_parallel_size):
+    print("Testing splitting size =", tensor_model_parallel_size)
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+    listy = []
+    for i in range(tensor_model_parallel_size):
+        listy.append(torch.randn(10, 1))
+    x = torch.cat(tuple(listy), 1)
+    out = mappings._split(x)
+    assert torch.equal(out, listy[parallel_state.get_tensor_model_parallel_rank()])
+    parallel_state.destroy_model_parallel()
+    print("Passed!")
+
+
+def test__gather(args, tensor_model_parallel_size):
+
+    print("Testing gathering size =", tensor_model_parallel_size)
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+    assert torch.equal(
+        mappings._gather(torch.tensor([parallel_state.get_tensor_model_parallel_rank()])),
+        torch.tensor(list(range(tensor_model_parallel_size))),
+    )
+    parallel_state.destroy_model_parallel()
+    print("Passed!")
+
+
+if __name__ == "__main__":
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    initialize_distributed()
+
+    world_size = torch.distributed.get_world_size()
+    args = global_vars.get_args()
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        test__reduce(args, tensor_model_parallel_size)
+        test__split(args, tensor_model_parallel_size)
+        test__gather(args, tensor_model_parallel_size)
+        tensor_model_parallel_size *= 2
+    print(">> passed the test :-)")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_megatron_gpt_pipeline.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_megatron_gpt_pipeline.py
new file mode 100644
index 0000000000..e48cbe42c0
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_megatron_gpt_pipeline.py
@@ -0,0 +1,132 @@
+from functools import partial
+from typing import List
+
+import torch
+
+from apex.transformer import parallel_state
+from apex.transformer.pipeline_parallel.schedules.common import _get_params_for_weight_decay_optimization
+from apex.transformer.pipeline_parallel.schedules.common import build_model
+from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_with_interleaving import _forward_backward_pipelining_with_interleaving
+from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_without_interleaving import forward_backward_pipelining_without_interleaving
+from apex.transformer.pipeline_parallel.utils import average_losses_across_data_parallel_group
+from apex.transformer.pipeline_parallel.utils import get_ltor_masks_and_position_ids
+from apex.transformer.pipeline_parallel.utils import setup_microbatch_calculator
+from apex.transformer.pipeline_parallel.utils import update_num_microbatches
+from apex.transformer.tensor_parallel import model_parallel_cuda_manual_seed
+from apex.transformer.testing import global_vars
+from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE
+from apex.transformer.testing.commons import initialize_distributed
+from apex.transformer.testing.commons import print_separator
+from apex.transformer.testing.standalone_gpt import gpt_model_provider
+from apex.transformer.utils import rank_print
+
+
+global_vars.set_global_variables()
+N_VOCAB = 8192
+
+
+def generate_batch(batch_size, sequence_length):
+    size = batch_size, sequence_length + 1
+    int_tensor = torch.randint(low=0, high=N_VOCAB, size=size, dtype=torch.long).cuda()
+    return int_tensor,
+
+
+# Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L44
+def get_batch(int_tensors: List[torch.Tensor]):
+    data = int_tensors[0]
+    # Unpack.
+    tokens_ = data.long()
+    labels = tokens_[:, 1:].contiguous()
+    tokens = tokens_[:, :-1].contiguous()
+    # Get the masks and position ids.
+    attention_mask, loss_mask, position_ids = get_ltor_masks_and_position_ids(
+        tokens,
+        N_VOCAB,  # tokenizer.eod,
+        False,  # args.reset_position_ids,
+        False,  # args.reset_attention_mask,
+        False,  # args.eod_mask_loss,
+    )
+    return tokens, labels, loss_mask, attention_mask, position_ids
+
+
+# Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L75
+def loss_func(loss_mask, output_tensor):
+    losses = output_tensor.float()
+    loss_mask = loss_mask.view(-1).float()
+    loss = torch.sum(losses.view(-1) * loss_mask) / loss_mask.sum()
+
+    # Reduce loss for logging.
+    averaged_loss = average_losses_across_data_parallel_group([loss])
+
+    return loss, {'lm loss': averaged_loss[0]}
+
+
+# Ref: https://github.com/NVIDIA/Megatron-LM/blob/b31e1296354e979722627a6c4dedafe19b51fa97/pretrain_gpt.py#L86
+# TODO (mkozuki): Currently I'm seeing no attribute `word_embeddings` which looks weird.
+def forward_step(batch, model):
+    """Forward step."""
+    tokens, labels, loss_mask, attention_mask, position_ids = get_batch(batch)
+    output_tensor = model(tokens, position_ids, attention_mask, labels=labels)
+    return output_tensor, partial(loss_func, loss_mask)
+
+
+def run_gpt(pipeline_model_parallel_size, virtual_pipeline_model_parallel_size=None, forward_only=False):
+    parallel_state.initialize_model_parallel(1, pipeline_model_parallel_size, virtual_pipeline_model_parallel_size)
+    model_parallel_cuda_manual_seed(42)
+
+    model = build_model(
+        gpt_model_provider, True,
+        virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size)
+    # rank_print("building model")
+    assert isinstance(model, list)
+    assert len(model) == (1 or virtual_pipeline_model_parallel_size)
+    _param_groups = _get_params_for_weight_decay_optimization(model)
+    torch.optim.Adam(_param_groups)
+
+    if parallel_state.is_pipeline_last_stage():
+        # rank_print("checking `word_embeddings` existence")
+        for m in model:
+            assert hasattr(m, "word_embeddings")
+
+    args = global_vars.get_args()
+    if virtual_pipeline_model_parallel_size is None:
+        batch = generate_batch(args.global_batch_size, args.seq_length)
+    else:
+        batch = [generate_batch(args.global_batch_size, args.seq_length) for _ in range(virtual_pipeline_model_parallel_size)]
+    # rank_print("preparing batch")
+
+    if virtual_pipeline_model_parallel_size is None:
+        fwd_bwd_func = forward_backward_pipelining_without_interleaving
+    else:
+        fwd_bwd_func = _forward_backward_pipelining_with_interleaving
+    # rank_print(f"selecting forward_backward func: {fwd_bwd_func}")
+
+    tensor_shape = (args.seq_length, args.micro_batch_size, args.hidden_size)
+    # rank_print(f"`tensor_shape`: {tensor_shape}")
+    fwd_bwd_func(forward_step, batch, model, forward_only=forward_only, tensor_shape=tensor_shape)
+
+    # rank_print(TEST_SUCCESS_MESSAGE)
+
+
+if __name__ == "__main__":
+    initialize_distributed()
+    args = global_vars.get_args()
+    args.padded_vocab_size = N_VOCAB
+    setup_microbatch_calculator(
+        args.rank,
+        args.rampup_batch_size,
+        args.global_batch_size,
+        args.micro_batch_size,
+        1,  # args.data_parallel_size,
+    )
+    update_num_microbatches(0, True)
+    print_separator("run GPT model")
+    try:
+        run_gpt(torch.distributed.get_world_size())
+    # TODO(mkozuki): handle exception correctly, but for now, lazily commenting out as
+    # this won't get kicked by CI
+    except Exception as e:
+        # rank_print(str(e))
+        pass
+    finally:
+        parallel_state.destroy_model_parallel()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_pipeline_parallel_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_pipeline_parallel_test.py
new file mode 100644
index 0000000000..c4824cfdd3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_pipeline_parallel_test.py
@@ -0,0 +1,198 @@
+from typing import Optional, Union, List
+
+import torch
+import torch.nn as nn
+
+from apex.transformer import parallel_state
+from apex.transformer.pipeline_parallel import get_forward_backward_func
+from apex.transformer.pipeline_parallel.schedules.common import _get_params_for_weight_decay_optimization
+from apex.transformer.pipeline_parallel.schedules.common import build_model
+from apex.transformer.pipeline_parallel.schedules.fwd_bwd_no_pipelining import forward_backward_no_pipelining
+from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_with_interleaving import _forward_backward_pipelining_with_interleaving
+from apex.transformer.pipeline_parallel.schedules.fwd_bwd_pipelining_without_interleaving import forward_backward_pipelining_without_interleaving
+from apex.transformer.pipeline_parallel.utils import average_losses_across_data_parallel_group
+from apex.transformer.pipeline_parallel.utils import setup_microbatch_calculator
+from apex.transformer.pipeline_parallel.utils import update_num_microbatches
+from apex.transformer.testing import global_vars
+from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE
+from apex.transformer.testing.commons import initialize_distributed
+from apex.transformer.testing.commons import print_separator
+from apex.transformer.utils import rank_print
+
+
+global_vars.set_global_variables()
+
+
+batch_size, micro_batch_size = None, None
+hidden_size = 16
+fwd_bwd_functions = {
+    "no_pipelining": forward_backward_no_pipelining,
+    "no_interleaving": forward_backward_pipelining_without_interleaving,
+    "interleaving": _forward_backward_pipelining_with_interleaving,
+}
+
+
+# note (mkozuki): `pre_process` and `post_process` are a placeholder until interleaving schedule test comes.
+class MyLayer(nn.Module):
+
+    def __init__(self, pre_process: bool, post_process: bool):
+        super().__init__()
+        self.pre_process = pre_process
+        self.post_process = post_process
+        self.layer = nn.Linear(hidden_size, hidden_size)
+
+    def forward(self, x):
+        return self.layer(x)
+
+class MyModel(nn.Module):
+
+    def __init__(self, pre_process: bool = False, post_process: bool = False) -> None:
+        super().__init__()
+        self.pre_process = pre_process
+        self.post_process = post_process
+        self.layer = MyLayer(pre_process=pre_process, post_process=post_process)
+        self.input_tensor = None
+
+    def set_input_tensor(self, input_tensor: Union[torch.Tensor, List[torch.Tensor]]) -> None:
+        self.input_tensor = input_tensor
+
+    def forward(self, x: Optional[torch.Tensor]) -> torch.Tensor:
+        if self.input_tensor is None:
+            return self.layer(x)
+        return self.layer(self.input_tensor)
+
+
+
+def model_provider_func(pre_process, post_process) -> MyModel:
+    return MyModel(pre_process, post_process)
+
+
+def process_batch(batch):
+    if isinstance(batch, list):
+        x = batch[0]
+    else:
+        x = batch
+    return x
+
+
+def fwd_step_func(batch, model):
+    x = process_batch(batch)
+    y = model(x)
+
+    # note (mkozuki): I don't think this function is nice but I do think this is enough for now
+    # just to check the sanity of ported pipeline functions.
+    def loss_func(x):
+        loss = torch.sum(x)
+        averaged_loss = average_losses_across_data_parallel_group([loss])
+        return loss, {'avg': averaged_loss}
+    return y, loss_func
+
+
+# TODO (mkozuki): Add a case with `autocast` and `GradScaler`.
+# Run forward & backward for one minibatch.
+def forward_backward_func_template(
+        name: str,
+        forward_backward_func,
+        pipeline_model_parallel_size: int,
+        forward_only: bool,
+) -> None:
+    print_separator(f"name: {name}, pipeline model parallel size: {pipeline_model_parallel_size}")
+    virtual_pipeline_model_parallel_size = 2 if name == "interleaving" else None
+    if name == "no_pipelining":
+        # note (mkozuki): `forward_backward_no_pipelining` is **NOTE** compatible with
+        # pipeline_model_parallel_size>1. So use pipeline_model_parallel_size as
+        # tensor_model_parallel_size and set pipeline_model_parallel_size to 1.
+        parallel_state.initialize_model_parallel(1, 1, None)
+    else:
+        # NOTE (mkozuki): `virtual_pipeline_model_parallel_size` is necessary to enable interleaving scheduling
+        # In megatron, `args.virtual_pipeline_model_parallel_size` is computed in megatron/arguments.py and
+        # used ubiquitously but this test uses custom model so it's safe to abuse.
+        parallel_state.initialize_model_parallel(
+            1, pipeline_model_parallel_size, virtual_pipeline_model_parallel_size)
+        if virtual_pipeline_model_parallel_size is not None:
+            # Check the experimental warning message
+            get_forward_backward_func(virtual_pipeline_model_parallel_size, pipeline_model_parallel_size)
+    pipeline_model_parallel_size = parallel_state.get_pipeline_model_parallel_world_size()
+
+    model = build_model(
+        model_provider_func,
+        wrap_with_ddp=True,
+        virtual_pipeline_model_parallel_size=virtual_pipeline_model_parallel_size,
+    )
+    assert isinstance(model, list)
+    assert len(model) == (1 if virtual_pipeline_model_parallel_size is None else virtual_pipeline_model_parallel_size)
+    _param_groups = _get_params_for_weight_decay_optimization(model)
+    torch.optim.Adam(_param_groups)
+
+    tensor_shape = [batch_size // parallel_state.get_data_parallel_world_size(), hidden_size]
+    batch = (torch.randn(tensor_shape).cuda(),)
+    tensor_shape[0] = micro_batch_size
+
+    update_num_microbatches(0)
+    forward_backward_func(
+        fwd_step_func, batch, model, forward_only=forward_only, tensor_shape=tensor_shape)
+
+    if not forward_only:
+        # rank_print("grad check")
+        for m in model:
+            for p in m.parameters():
+                if p.grad is None:
+                    raise RuntimeError("grad not found")
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(TEST_SUCCESS_MESSAGE)
+
+
+if __name__ == "__main__":
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    n_tests = 0
+    failures = []
+
+    initialize_distributed()
+    world_size = torch.distributed.get_world_size()
+    args = global_vars.get_args()
+    batch_size = args.global_batch_size
+    micro_batch_size = args.micro_batch_size
+    setup_microbatch_calculator(
+        args.rank,
+        args.rampup_batch_size,
+        args.global_batch_size,
+        args.micro_batch_size,
+        1,  # args.data_parallel_size,
+    )
+    for forward_only in (True, False):
+        for name, forward_backward_func in fwd_bwd_functions.items():
+            n_tests += 1
+            # TODO (mkozuki): Test with data parallel size > 1.
+            pipeline_model_parallel_size = world_size
+            try:
+                forward_backward_func_template(
+                    name,
+                    forward_backward_func,
+                    pipeline_model_parallel_size,
+                    forward_only,
+                )
+            except Exception as e:
+                failures.append(
+                    f"\t# {name} failed with pipeline size: {pipeline_model_parallel_size} "
+                    f"and forward_only: {forward_only}\n"
+                    f"pipeline rank: {parallel_state.get_pipeline_model_parallel_rank()}, "
+                    f"virtual pipeline rank: {parallel_state.get_virtual_pipeline_model_parallel_rank()}\n"
+                    f"{str(e)}"
+                )
+            finally:
+                parallel_state.destroy_model_parallel()
+        else:
+            print_separator(f"{name} works")
+    print_separator("TEST RESULT")
+    if failures:
+        torch.distributed.barrier()
+        if torch.distributed.get_rank() == 0:
+            print("\n".join(failures))
+        msg = f"{len(failures)} / {n_tests} cases failed"
+        raise RuntimeError(msg)
+    else:
+        torch.distributed.barrier()
+        if torch.distributed.get_rank() == 0:
+            print("### PASS!")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_random_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_random_test.py
new file mode 100644
index 0000000000..c25e55319c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_random_test.py
@@ -0,0 +1,213 @@
+# coding=utf-8
+# Copyright (c) 2021, NVIDIA CORPORATION.  All rights reserved.
+#
+# 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.
+import torch
+
+from apex.transformer import parallel_state
+from apex.transformer import tensor_parallel
+from apex.transformer.testing import global_vars
+from apex.transformer.testing.commons import print_separator
+from apex.transformer.testing.commons import initialize_distributed
+from apex.transformer.testing.commons import TEST_SUCCESS_MESSAGE
+
+
+global_vars.set_global_variables()
+
+
+def test_set_cuda_rng_state(tensor_model_parallel_size):
+
+    if torch.distributed.get_rank() == 0:
+        print('> testing set_rng_state with size {} ...'.
+              format(tensor_model_parallel_size))
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    size = 123
+    seed = 1234
+    torch.cuda.manual_seed(seed)
+    tensor = torch.cuda.FloatTensor(size)
+
+    # Get the state
+    rng_state = torch.cuda.get_rng_state()
+    rng_state_copy = rng_state.clone()
+
+    # Do some stuff.
+    for _ in range(5):
+        torch.randn(size, out=tensor)
+    result_1 = tensor.clone()
+
+    assert rng_state.sub(rng_state_copy).max() == 0
+    assert torch.cuda.get_rng_state().sub(rng_state_copy).max() > 0
+
+    # State should be different.
+    new_rng_state = torch.cuda.get_rng_state()
+    max_diff = new_rng_state.sub(rng_state).max()
+    print('   max diff in rng state (should be non-zero) on global rank {}: {}'.
+          format(torch.distributed.get_rank(), max_diff))
+    assert max_diff > 0
+
+    # Reset the rng state and do the same stuff.
+    tensor_parallel.random._set_cuda_rng_state(rng_state)
+    for _ in range(5):
+        torch.randn(size, out=tensor)
+    tensor_parallel.random._set_cuda_rng_state(rng_state)
+    for _ in range(5):
+        torch.randn(size, out=tensor)
+    result_2 = tensor.clone()
+
+    # Results should be the same
+    error = result_2.sub(result_1).abs().max()
+    print('   max error in generated tensors (should be zero) on '
+          'global rank {}: {}'.format(torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    # Input state should have remained intact.
+    error = rng_state.sub(rng_state_copy).max()
+    print('   max error in rng state (should be zero) on global rank {}: {}'.
+          format(torch.distributed.get_rank(), error))
+    assert error == 0
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(TEST_SUCCESS_MESSAGE)
+
+
+def test_cuda_rng_tracker(tensor_model_parallel_size):
+
+    if torch.distributed.get_rank() == 0:
+        print('> testing cuda rng tracker with size {} ...'.
+              format(tensor_model_parallel_size))
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    seed_1 = 1234
+    seed_2 = 4321
+    size = [12, 21]
+    tensor = torch.cuda.FloatTensor(size)
+
+    # Set to seed_1 and generate two tensors.
+    torch.cuda.manual_seed(seed_1)
+    torch.randn(size, out=tensor)
+    target_11 = tensor.clone()
+    torch.randn(size, out=tensor)
+    target_12 = tensor.clone()
+
+    # Set to seed_2 and generate two tensors.
+    torch.cuda.manual_seed(seed_2)
+    torch.randn(size, out=tensor)
+    target_21 = tensor.clone()
+    torch.randn(size, out=tensor)
+    target_22 = tensor.clone()
+
+    # Now if we interleave seed_1 and seed_2,
+    # we should still get the same tensors
+    torch.cuda.manual_seed(seed_1)
+    tensor_parallel.random.get_cuda_rng_tracker().add('test', seed_2)
+
+    torch.randn(size, out=tensor)
+    result_11 = tensor.clone()
+
+    with tensor_parallel.random.get_cuda_rng_tracker().fork('test'):
+        torch.randn(size, out=tensor)
+        result_21 = tensor.clone()
+
+    torch.randn(size, out=tensor)
+    result_12 = tensor.clone()
+
+    with tensor_parallel.random.get_cuda_rng_tracker().fork('test'):
+        torch.randn(size, out=tensor)
+        result_22 = tensor.clone()
+
+    diff = result_11.sub(result_21).abs().max()
+    diff = min(diff, result_12.sub(result_22).abs().max())
+    print('   max diff in generated tensors (should be non-zero) on '
+          'global rank {}: {}'.format(torch.distributed.get_rank(), diff))
+    assert diff > 1.0e-6
+    error = max(result_11.sub(target_11).abs().max(),
+                result_12.sub(target_12).abs().max())
+    error = max(error, result_21.sub(target_21).abs().max())
+    error = max(error, result_22.sub(target_22).abs().max())
+    print('   max error in generated tensors (should be zero) on '
+          'global rank {}: {}'.format(torch.distributed.get_rank(), error))
+    assert error < 1.0e-6
+
+    # Reset the tracker
+    tensor_parallel.random.get_cuda_rng_tracker().reset()
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(TEST_SUCCESS_MESSAGE)
+
+
+def test_model_parallel_cuda_manual_seed(tensor_model_parallel_size):
+
+    if torch.distributed.get_rank() == 0:
+        print(
+            '> testing model parallel cuda manual seed with size {} ...'.format(
+                tensor_model_parallel_size))
+
+    parallel_state.initialize_model_parallel(tensor_model_parallel_size)
+    tensor_model_parallel_size = parallel_state.get_tensor_model_parallel_world_size()
+
+    tensor_parallel.random.model_parallel_cuda_manual_seed(12345)
+    assert torch.cuda.initial_seed() == 12345
+    with tensor_parallel.random.get_cuda_rng_tracker().fork():
+        assert (
+            torch.cuda.initial_seed() ==
+            12345 + 2718 + parallel_state.get_tensor_model_parallel_rank()
+        )
+
+    # Reset the tracker
+    tensor_parallel.random.get_cuda_rng_tracker().reset()
+
+    # Reset groups
+    parallel_state.destroy_model_parallel()
+
+    torch.distributed.barrier()
+    if torch.distributed.get_rank() == 0:
+        print(TEST_SUCCESS_MESSAGE)
+
+
+if __name__ == '__main__':
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+
+    initialize_distributed()
+    world_size = torch.distributed.get_world_size()
+
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        print_separator('test set rng state')
+        test_set_cuda_rng_state(tensor_model_parallel_size)
+        tensor_model_parallel_size *= 2
+
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        print_separator('test cuda rng tracker')
+        test_cuda_rng_tracker(tensor_model_parallel_size)
+        tensor_model_parallel_size *= 2
+
+    tensor_model_parallel_size = 1
+    while tensor_model_parallel_size <= world_size:
+        print_separator('test model parallel cuda manual seed')
+        test_model_parallel_cuda_manual_seed(tensor_model_parallel_size)
+        tensor_model_parallel_size *= 2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_utils_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_utils_test.py
new file mode 100644
index 0000000000..654dfe23ac
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/run_utils_test.py
@@ -0,0 +1,22 @@
+import torch
+
+from apex.transformer.tensor_parallel import utils
+
+
+def test_divide():
+    assert utils.divide(8, 4) == 2
+
+
+def test_split_tensor_along_last_dim():
+    inputy = torch.randn((100, 100, 100))
+    splits = utils.split_tensor_along_last_dim(inputy, 10)
+    last_dim_shapes = torch.tensor([int(split.size()[-1]) for split in splits])
+    assert torch.equal(last_dim_shapes, torch.full((10,), 10))
+
+
+if __name__ == "__main__":
+    torch.backends.cuda.matmul.allow_tf32 = False
+    torch.backends.cudnn.allow_tf32 = False
+    test_divide()
+    test_split_tensor_along_last_dim()
+    print(">> passed the test :-)")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_batch_sampler.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_batch_sampler.py
new file mode 100644
index 0000000000..0818d906bc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_batch_sampler.py
@@ -0,0 +1,149 @@
+from itertools import product
+import unittest
+
+import torch
+from torch.utils.data import Dataset
+from torch.utils.data import RandomSampler
+from torch.utils.data import BatchSampler
+from torch.utils.data import DataLoader
+
+from apex.transformer.pipeline_parallel.utils import _split_batch_into_microbatch as split_batch_into_microbatch
+
+
+class MyIterableDataset(Dataset):
+    def __init__(self, start, end):
+        super().__init__()
+        assert end > start, "this example code only works with end >= start"
+        self.start = start
+        self.end = end
+        self.samples = list(range(self.start, self.end))
+
+    def __iter__(self):
+        return iter(range(self.start, self.end))
+
+    def __getitem__(self, index):
+        return self.samples[index]
+
+
+class MegatronPretrainingRandomSampler:
+
+    def __init__(self, total_samples, consumed_samples, micro_batch_size,
+                 data_parallel_rank, data_parallel_size):
+        # Keep a copy of input params for later use.
+        self.total_samples = total_samples
+        self.consumed_samples = consumed_samples
+        self.micro_batch_size = micro_batch_size
+        self.data_parallel_rank = data_parallel_rank
+        self.data_parallel_size = data_parallel_size
+        self.micro_batch_times_data_parallel_size = \
+            self.micro_batch_size * data_parallel_size
+        self.last_batch_size = \
+            self.total_samples % self.micro_batch_times_data_parallel_size
+
+        # Sanity checks.
+        assert self.total_samples > 0, \
+            'no sample to consume: {}'.format(self.total_samples)
+        assert self.micro_batch_size > 0
+        assert data_parallel_size > 0
+        assert self.data_parallel_rank < data_parallel_size, \
+            'data_parallel_rank should be smaller than data size: {}, ' \
+            '{}'.format(self.data_parallel_rank, data_parallel_size)
+
+    def __len__(self):
+        return self.total_samples
+
+    def __iter__(self):
+        active_total_samples = self.total_samples - self.last_batch_size
+        self.epoch = self.consumed_samples // active_total_samples
+        current_epoch_samples = self.consumed_samples % active_total_samples
+        assert current_epoch_samples % self.micro_batch_times_data_parallel_size == 0
+
+        # data sharding and random sampling
+        bucket_size = (self.total_samples // self.micro_batch_times_data_parallel_size) * self.micro_batch_size
+        bucket_offset = current_epoch_samples // self.data_parallel_size
+        start_idx = self.data_parallel_rank * bucket_size
+
+        g = torch.Generator()
+        g.manual_seed(self.epoch)
+        random_idx = torch.randperm(bucket_size, generator=g).tolist()
+        idx_range = [start_idx + x for x in random_idx[bucket_offset:]]
+
+        batch = []
+        # Last batch if not complete will be dropped.
+        for idx in idx_range:
+            batch.append(idx)
+            if len(batch) == self.micro_batch_size:
+                self.consumed_samples += self.micro_batch_times_data_parallel_size
+                yield batch
+                batch = []
+
+
+# Samples 8 tensors in total.
+# First sample 4 tensors twice, then sample 2 tensors fourth.
+class TestBatchSamplerBehavior(unittest.TestCase):
+    def test_batch_sampler_behavior(self):
+        dataset = MyIterableDataset(0, 100)
+
+        for num_workers in (1, 2, 4):
+            with self.subTest(f"{num_workers}"):
+                torch.manual_seed(42)
+                loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, 4, 0, 1), num_workers=num_workers)
+                samples = []
+                for i, batch in enumerate(loader):
+                    samples.append(batch)
+                    if i == 2 - 1:
+                        break
+
+                torch.manual_seed(42)
+                loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, 2, 0, 1), num_workers=num_workers)
+                samples2 = []
+                for i, batch in enumerate(loader):
+                    samples2.append(batch)
+                    if i == 4 - 1:
+                        break
+                torch.testing.assert_allclose(torch.cat(samples), torch.cat(samples2))
+
+    def test_split_batch(self):
+
+        class MyIterableDataset(Dataset):
+            def __init__(self, start, end):
+                super().__init__()
+                assert end > start, "this example code only works with end >= start"
+                self.start = start
+                self.end = end
+                self.samples = list(range(self.start, self.end))
+
+            def __len__(self):
+                return self.end - self.start
+
+            def __iter__(self):
+                return iter(range(self.start, self.end))
+
+            def __getitem__(self, index):
+                return (torch.tensor([index, index]), torch.tensor([index // 2, index // 2]))
+
+        dataset = MyIterableDataset(0, 100)
+        torch.manual_seed(42)
+        global_batch_size = 16
+        loader = DataLoader(dataset, batch_sampler=MegatronPretrainingRandomSampler(100, 0, global_batch_size, 0, 1), num_workers=2)
+        batch = next(iter(loader))
+        # samples = None
+        # for i, batch in enumerate(loader):
+        #     # samples = batch
+        #     if i == 0:
+        #         break
+
+        for _micro_batch_size in (1, 2, 4, 8):
+            microbatches = list(split_batch_into_microbatch(
+                batch,
+                _micro_batch_size=_micro_batch_size,
+                _global_batch_size=global_batch_size,
+            ))
+            # print(batch)
+            # print(microbatches)
+            self.assertEqual(len(microbatches), global_batch_size // _micro_batch_size)
+            self.assertEqual(len(microbatches[0][0]), _micro_batch_size)
+
+
+if __name__ == "__main__":
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_fused_softmax.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_fused_softmax.py
new file mode 100644
index 0000000000..fb6a9c5537
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_fused_softmax.py
@@ -0,0 +1,168 @@
+"""Test for fused softmax functions.
+
+Ref: https://github.com/NVIDIA/Megatron-LM/blob/40becfc96c4144985458ac0e0fae45dbb111fbd2/megatron/fused_kernels/tests/test_fused_kernels.py
+"""  # NOQA
+import itertools
+import unittest
+
+import torch
+
+from apex.transformer import AttnMaskType
+from apex.transformer.functional import FusedScaleMaskSoftmax
+
+
+def attention_mask_func(attention_scores, attention_mask):
+    return attention_scores.masked_fill(attention_mask, -10000.0)
+
+
+autocast_dtypes = (torch.half, torch.bfloat16) if torch.cuda.is_bf16_supported() else (torch.half,)
+
+
+class TestFusedScaleMaskSoftmax(unittest.TestCase):
+
+    def _setup_fused_softmax(self, input_in_fp16, input_in_bf16, scale=None, softmax_in_fp32=False, attn_mask_type=AttnMaskType.padding):
+        fused_fn = FusedScaleMaskSoftmax(
+            input_in_fp16=input_in_fp16,
+            input_in_bf16=input_in_bf16,
+            mask_func=attention_mask_func,
+            scale=scale,
+            softmax_in_fp32=softmax_in_fp32,
+            attn_mask_type=attn_mask_type,
+            scaled_masked_softmax_fusion=True,
+        )
+        torch_fn = FusedScaleMaskSoftmax(
+            input_in_fp16=input_in_fp16,
+            input_in_bf16=input_in_bf16,
+            mask_func=attention_mask_func,
+            scale=scale,
+            softmax_in_fp32=softmax_in_fp32,
+            attn_mask_type=attn_mask_type,
+            scaled_masked_softmax_fusion=False,
+        )
+        return fused_fn, torch_fn
+
+    def test_fused_scale_mask_softmax(self):
+        """
+        attention_scores.shape = [4, 12, 24, 24]
+        mask.shape = [4, 1, 24, 24]
+        """
+        for (dtype, scale, softmax_in_fp32) in itertools.product(
+                (torch.half, torch.bfloat16),
+                (None, 2.0),
+                (False, True),
+        ):
+            with self.subTest(f"{dtype}-{scale}-{softmax_in_fp32}"):
+                input_in_fp16 = dtype == torch.half
+                input_in_bf16 = dtype == torch.bfloat16
+                if not (scale is None or softmax_in_fp32):
+                    with self.assertRaises(RuntimeError):
+                        self._setup_fused_softmax(input_in_fp16, input_in_bf16, scale, softmax_in_fp32, AttnMaskType.padding)
+                    return
+                fused_fn, torch_fn = self._setup_fused_softmax(input_in_fp16, input_in_bf16, scale, softmax_in_fp32, AttnMaskType.padding)
+
+                attention_scores_0 = torch.randn((4, 12, 24, 24)).to(device="cuda", dtype=dtype).requires_grad_(True)
+                with torch.no_grad():
+                    attention_scores_1 = attention_scores_0.clone().requires_grad_(True)
+                mask = torch.randint(0, 2, (4, 1, 24, 24), device="cuda").bool()
+                expected = fused_fn(attention_scores_0, mask)
+                actual = torch_fn(attention_scores_1, mask)
+                torch.testing.assert_allclose(actual, expected)
+
+                g0 = torch.rand_like(actual)
+                with torch.no_grad():
+                    g1 = g0.clone()
+                expected.backward(g0)
+                actual.backward(g1)
+
+    def test_autocast_fused_scale_mask_softmax(self):
+        for dtype in autocast_dtypes:
+            with self.subTest(f"{dtype}"):
+                input_in_fp16 = dtype == torch.half
+                input_in_bf16 = dtype == torch.bfloat16
+                fused_fn, torch_fn = self._setup_fused_softmax(
+                    input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.padding)
+
+                attention_scores_0 = torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True)
+                with torch.no_grad():
+                    attention_scores_1 = attention_scores_0.clone().to(dtype).requires_grad_(True)
+                mask = torch.randint(0, 2, (4, 1, 24, 24)).bool().cuda()
+
+                expected = torch_fn(attention_scores_1, mask)
+                with torch.cuda.amp.autocast(dtype=dtype):
+                    actual = fused_fn(attention_scores_0, mask)
+                    self.assertEqual(actual.dtype, dtype)
+                torch.testing.assert_allclose(actual, expected)
+
+                g0 = torch.rand_like(actual)
+                with torch.no_grad():
+                    g1 = g0.clone()
+                expected.backward(g0)
+                actual.backward(g1)
+
+    def test_fused_upper_triangle_mask_softmax(self):
+        """
+        attn_weights.shape: [4, 12, 24, 24]
+        total_mask.shape: [4, 1, 24, 24]
+
+        total_mask[0, 0], a 24x24 matrix is like a lower triangular matrix, but
+        upper elements are True and lower elements and diagonal are False.
+        """
+        for (dtype, scale, softmax_in_fp32) in itertools.product(
+                (torch.half, torch.bfloat16),
+                (None, 2.0),
+                (False, True),
+        ):
+            with self.subTest(f"{dtype}-{scale}-{softmax_in_fp32}"):
+                input_in_fp16 = dtype == torch.half
+                input_in_bf16 = dtype == torch.bfloat16
+                if not (scale is None or softmax_in_fp32):
+                    with self.assertRaises(RuntimeError):
+                        self._setup_fused_softmax(
+                            input_in_fp16, input_in_bf16, scale, softmax_in_fp32, AttnMaskType.causal)
+                    return
+                fused_fn, torch_fn = self._setup_fused_softmax(
+                    input_in_fp16, input_in_bf16, scale, softmax_in_fp32, AttnMaskType.causal)
+
+                attn_weights_0 = torch.randn((4, 12, 24, 24)).to(device="cuda", dtype=dtype).requires_grad_(True)
+                with torch.no_grad():
+                    attn_weights_1 = attn_weights_0.clone().requires_grad_(True)
+                total_mask = (~(
+                    torch.tril(torch.randn((24, 24), device="cuda")).bool()
+                ).unsqueeze(0).unsqueeze(0))
+                total_mask = total_mask.repeat((4, 1, 1, 1))
+                expected = fused_fn(attn_weights_0, total_mask)
+                actual = torch_fn(attn_weights_1, total_mask)
+                torch.testing.assert_allclose(actual, expected)
+
+                g0 = torch.randn_like(actual)
+                with torch.no_grad():
+                    g1 = g0.clone()
+                actual.backward(g0)
+                expected.backward(g1)
+
+    def test_autocast_fused_upper_triangle_mask_softmax(self):
+        for dtype in autocast_dtypes:
+            with self.subTest(f"{dtype}"):
+                input_in_fp16 = dtype == torch.half
+                input_in_bf16 = dtype == torch.bfloat16
+                fused_fn, torch_fn = self._setup_fused_softmax(
+                    input_in_fp16, input_in_bf16, attn_mask_type=AttnMaskType.causal)
+
+                attn_weights_0 = torch.randn((4, 12, 24, 24)).cuda().requires_grad_(True)
+                with torch.no_grad():
+                    attn_weights_1 = attn_weights_0.clone().to(dtype).requires_grad_(True)
+                total_mask = (~(
+                    torch.tril(torch.randn((24, 24), device="cuda")).bool()
+                ).unsqueeze(0).unsqueeze(0))
+
+                with torch.cuda.amp.autocast(dtype=dtype):
+                    actual = fused_fn(attn_weights_0, total_mask)
+                    self.assertEqual(actual.dtype, dtype)
+                expected = torch_fn(attn_weights_1, total_mask)
+                torch.testing.assert_allclose(actual, expected)
+
+                g0 = torch.randn_like(actual)
+                with torch.no_grad():
+                    g1 = g0.clone()
+                actual.backward(g0)
+                expected.backward(g1)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_transformer_module.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_transformer_module.py
new file mode 100644
index 0000000000..b9d87c80a3
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L0/run_transformer/test_transformer_module.py
@@ -0,0 +1,89 @@
+from typing import Tuple
+import os
+import subprocess
+import sys
+import unittest
+
+
+DENY_TEST = [
+    "megatron_gpt_pipeline",
+]
+MULTIGPU_TEST = [
+    "pipeline_parallel_test",
+]
+
+
+def get_launch_option(test_filename) -> Tuple[bool, str]:
+    should_skip = False
+    for multigpu_test in MULTIGPU_TEST:
+        if multigpu_test in test_filename:
+            import torch
+            num_devices = torch.cuda.device_count()
+            if num_devices < 2:
+                should_skip = True
+            distributed_run_options = f"-m torch.distributed.run --nproc_per_node={num_devices}"
+            return should_skip, distributed_run_options
+    return should_skip, ""
+
+
+def run_transformer_tests():
+    python_executable_path = sys.executable
+    # repository_root = os.path.join(os.path.dirname(__file__), "../../../")
+    # directory = os.path.abspath(os.path.join(repository_root, "tests/mpu"))
+    directory = os.path.dirname(__file__)
+    files = [
+        os.path.join(directory, f) for f in os.listdir(directory)
+        if f.startswith("run_") and os.path.isfile(os.path.join(directory, f))
+    ]
+    print("#######################################################")
+    print(f"# Python executable path: {python_executable_path}")
+    print(f"# {len(files)} tests: {files}")
+    print("#######################################################")
+    errors = []
+    for i, test_file in enumerate(files, 1):
+        is_denied = False
+        for deny_file in DENY_TEST:
+            if deny_file in test_file:
+                is_denied = True
+        if is_denied:
+            print(f"### {i} / {len(files)}: {test_file} skipped")
+            continue
+        should_skip, launch_option = get_launch_option(test_file)
+        if should_skip:
+            print(f"### {i} / {len(files)}: {test_file} skipped. Requires multiple GPUs.")
+            continue
+        test_run_cmd = (
+            f"{python_executable_path} {launch_option} {test_file} "
+            "--micro-batch-size 2 --num-layers 1 --hidden-size 256 --num-attention-heads 8 --max-position-embeddings "
+            "32 --encoder-seq-length 32 --use-cpu-initialization"
+        )
+        print(f"### {i} / {len(files)}: cmd: {test_run_cmd}")
+        try:
+            output = subprocess.check_output(
+                test_run_cmd, shell=True
+            ).decode(sys.stdout.encoding).strip()
+        except Exception as e:
+            errors.append((test_file, str(e)))
+        else:
+            if '>> passed the test :-)' not in output:
+                errors.append((test_file, output))
+    else:
+        if not errors:
+            print("### PASSED")
+        else:
+            print("### FAILED")
+            short_msg = f"{len(errors)} out of {len(files)} tests failed"
+            print(short_msg)
+            for (filename, log) in errors:
+                print(f"File: {filename}\nLog: {log}")
+            raise RuntimeError(short_msg)
+
+
+class TestTransformer(unittest.TestCase):
+
+    def test_transformer(self):
+        run_transformer_tests()
+
+
+if __name__ == '__main__':
+    unittest.main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/compare.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/compare.py
new file mode 100644
index 0000000000..74374d412d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/compare.py
@@ -0,0 +1,64 @@
+import argparse
+import torch
+
+parser = argparse.ArgumentParser(description='Compare')
+parser.add_argument('--opt-level', type=str)
+parser.add_argument('--keep-batchnorm-fp32', type=str, default=None)
+parser.add_argument('--loss-scale', type=str, default=None)
+parser.add_argument('--fused-adam', action='store_true')
+parser.add_argument('--use_baseline', action='store_true')
+args = parser.parse_args()
+
+base_file = str(args.opt_level) + "_" +\
+            str(args.loss_scale) + "_" +\
+            str(args.keep_batchnorm_fp32) + "_" +\
+            str(args.fused_adam)
+
+file_e = "True_" + base_file
+file_p = "False_" + base_file
+if args.use_baseline:
+    file_b = "baselines/True_" + base_file
+
+dict_e = torch.load(file_e)
+dict_p = torch.load(file_p)
+if args.use_baseline:
+    dict_b = torch.load(file_b)
+
+torch.set_printoptions(precision=10)
+
+print(file_e)
+print(file_p)
+if args.use_baseline:
+    print(file_b)
+
+# ugly duplication here...
+if not args.use_baseline:
+    for n, (i_e, i_p) in enumerate(zip(dict_e["Iteration"], dict_p["Iteration"])):
+        assert i_e == i_p, "i_e = {}, i_p = {}".format(i_e, i_p)
+
+        loss_e = dict_e["Loss"][n]
+        loss_p = dict_p["Loss"][n]
+        assert loss_e == loss_p, "Iteration {}, loss_e = {}, loss_p = {}".format(i_e, loss_e, loss_p)
+        print("{:4} {:15.10f} {:15.10f} {:15.10f} {:15.10f}".format(
+              i_e,
+              loss_e,
+              loss_p,
+              dict_e["Speed"][n],
+              dict_p["Speed"][n]))
+else:
+    for n, (i_e, i_p) in enumerate(zip(dict_e["Iteration"], dict_p["Iteration"])):
+        assert i_e == i_p, "i_e = {}, i_p = {}".format(i_e, i_p)
+
+        loss_e = dict_e["Loss"][n]
+        loss_p = dict_p["Loss"][n]
+        loss_b = dict_b["Loss"][n]
+        assert loss_e == loss_p, "Iteration {}, loss_e = {}, loss_p = {}".format(i_e, loss_e, loss_p)
+        assert loss_e == loss_b, "Iteration {}, loss_e = {}, loss_b = {}".format(i_e, loss_e, loss_b)
+        print("{:4} {:15.10f} {:15.10f} {:15.10f} {:15.10f} {:15.10f} {:15.10f}".format(
+              i_e,
+              loss_b,
+              loss_e,
+              loss_p,
+              dict_b["Speed"][n],
+              dict_e["Speed"][n],
+              dict_p["Speed"][n]))
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/main_amp.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/main_amp.py
new file mode 100644
index 0000000000..106a0f6378
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/main_amp.py
@@ -0,0 +1,526 @@
+import argparse
+import os
+import shutil
+import time
+
+import torch
+import torch.nn as nn
+import torch.nn.parallel
+import torch.backends.cudnn as cudnn
+import torch.distributed as dist
+import torch.optim
+import torch.utils.data
+import torch.utils.data.distributed
+import torchvision.transforms as transforms
+import torchvision.datasets as datasets
+import torchvision.models as models
+
+import numpy as np
+
+try:
+    from apex.parallel import DistributedDataParallel as DDP
+    from apex.fp16_utils import *
+    from apex import amp, optimizers
+    from apex.multi_tensor_apply import multi_tensor_applier
+except ImportError:
+    raise ImportError("Please install apex from https://www.github.com/nvidia/apex to run this example.")
+
+model_names = sorted(name for name in models.__dict__
+                     if name.islower() and not name.startswith("__")
+                     and callable(models.__dict__[name]))
+
+parser = argparse.ArgumentParser(description='PyTorch ImageNet Training')
+parser.add_argument('data', metavar='DIR',
+                    help='path to dataset')
+parser.add_argument('--arch', '-a', metavar='ARCH', default='resnet18',
+                    choices=model_names,
+                    help='model architecture: ' +
+                    ' | '.join(model_names) +
+                    ' (default: resnet18)')
+parser.add_argument('-j', '--workers', default=4, type=int, metavar='N',
+                    help='number of data loading workers (default: 4)')
+parser.add_argument('--epochs', default=90, type=int, metavar='N',
+                    help='number of total epochs to run')
+parser.add_argument('--start-epoch', default=0, type=int, metavar='N',
+                    help='manual epoch number (useful on restarts)')
+parser.add_argument('-b', '--batch-size', default=256, type=int,
+                    metavar='N', help='mini-batch size per process (default: 256)')
+parser.add_argument('--lr', '--learning-rate', default=0.1, type=float,
+                    metavar='LR', help='Initial learning rate.  Will be scaled by <global batch size>/256: args.lr = args.lr*float(args.batch_size*args.world_size)/256.  A warmup schedule will also be applied over the first 5 epochs.')
+parser.add_argument('--momentum', default=0.9, type=float, metavar='M',
+                    help='momentum')
+parser.add_argument('--weight-decay', '--wd', default=1e-4, type=float,
+                    metavar='W', help='weight decay (default: 1e-4)')
+parser.add_argument('--print-freq', '-p', default=10, type=int,
+                    metavar='N', help='print frequency (default: 10)')
+parser.add_argument('--resume', default='', type=str, metavar='PATH',
+                    help='path to latest checkpoint (default: none)')
+parser.add_argument('-e', '--evaluate', dest='evaluate', action='store_true',
+                    help='evaluate model on validation set')
+parser.add_argument('--pretrained', dest='pretrained', action='store_true',
+                    help='use pre-trained model')
+
+parser.add_argument('--prof', dest='prof', action='store_true',
+                    help='Only run 10 iterations for profiling.')
+parser.add_argument('--deterministic', action='store_true')
+
+parser.add_argument("--local_rank", default=0, type=int)
+parser.add_argument('--sync_bn', action='store_true',
+                    help='enabling apex sync BN.')
+
+parser.add_argument('--has-ext', action='store_true')
+parser.add_argument('--opt-level', type=str)
+parser.add_argument('--keep-batchnorm-fp32', type=str, default=None)
+parser.add_argument('--loss-scale', type=str, default=None)
+parser.add_argument('--fused-adam', action='store_true')
+
+parser.add_argument('--prints-to-process', type=int, default=10)
+
+cudnn.benchmark = True
+
+def fast_collate(batch):
+    imgs = [img[0] for img in batch]
+    targets = torch.tensor([target[1] for target in batch], dtype=torch.int64)
+    w = imgs[0].size[0]
+    h = imgs[0].size[1]
+    tensor = torch.zeros( (len(imgs), 3, h, w), dtype=torch.uint8 )
+    for i, img in enumerate(imgs):
+        nump_array = np.asarray(img, dtype=np.uint8)
+        if(nump_array.ndim < 3):
+            nump_array = np.expand_dims(nump_array, axis=-1)
+        nump_array = np.rollaxis(nump_array, 2)
+
+        tensor[i] += torch.from_numpy(nump_array)
+        
+    return tensor, targets
+
+best_prec1 = 0
+args = parser.parse_args()
+
+# Let multi_tensor_applier be the canary in the coalmine
+# that verifies if the backend is what we think it is
+assert multi_tensor_applier.available == args.has_ext 
+
+print("opt_level = {}".format(args.opt_level))
+print("keep_batchnorm_fp32 = {}".format(args.keep_batchnorm_fp32), type(args.keep_batchnorm_fp32))
+print("loss_scale = {}".format(args.loss_scale), type(args.loss_scale))
+
+
+print("\nCUDNN VERSION: {}\n".format(torch.backends.cudnn.version()))
+
+if args.deterministic:
+    cudnn.benchmark = False
+    cudnn.deterministic = True
+    torch.manual_seed(args.local_rank)
+    torch.set_printoptions(precision=10)
+
+def main():
+    global best_prec1, args
+
+    args.distributed = False
+    if 'WORLD_SIZE' in os.environ:
+        args.distributed = int(os.environ['WORLD_SIZE']) > 1
+
+    args.gpu = 0
+    args.world_size = 1
+
+    if args.distributed:
+        args.gpu = args.local_rank % torch.cuda.device_count()
+        torch.cuda.set_device(args.gpu)
+        torch.distributed.init_process_group(backend='nccl',
+                                             init_method='env://')
+        args.world_size = torch.distributed.get_world_size()
+
+    assert torch.backends.cudnn.enabled, "Amp requires cudnn backend to be enabled."
+
+    # create model
+    if args.pretrained:
+        print("=> using pre-trained model '{}'".format(args.arch))
+        model = models.__dict__[args.arch](pretrained=True)
+    else:
+        print("=> creating model '{}'".format(args.arch))
+        model = models.__dict__[args.arch]()
+
+    if args.sync_bn:
+        import apex
+        print("using apex synced BN")
+        model = apex.parallel.convert_syncbn_model(model)
+
+    model = model.cuda()
+
+    # Scale learning rate based on global batch size
+    args.lr = args.lr*float(args.batch_size*args.world_size)/256. 
+    if args.fused_adam:
+        optimizer = optimizers.FusedAdam(model.parameters())
+    else:
+        optimizer = torch.optim.SGD(model.parameters(), args.lr,
+                                    momentum=args.momentum,
+                                    weight_decay=args.weight_decay)
+
+    model, optimizer = amp.initialize(
+        model, optimizer,
+        # enabled=False,
+        opt_level=args.opt_level,
+        keep_batchnorm_fp32=args.keep_batchnorm_fp32,
+        loss_scale=args.loss_scale
+        )
+
+    if args.distributed:
+        # By default, apex.parallel.DistributedDataParallel overlaps communication with 
+        # computation in the backward pass.
+        # model = DDP(model)
+        # delay_allreduce delays all communication to the end of the backward pass.
+        model = DDP(model, delay_allreduce=True)
+
+    # define loss function (criterion) and optimizer
+    criterion = nn.CrossEntropyLoss().cuda()
+
+    # Optionally resume from a checkpoint
+    if args.resume:
+        # Use a local scope to avoid dangling references
+        def resume():
+            if os.path.isfile(args.resume):
+                print("=> loading checkpoint '{}'".format(args.resume))
+                checkpoint = torch.load(args.resume, map_location = lambda storage, loc: storage.cuda(args.gpu))
+                args.start_epoch = checkpoint['epoch']
+                best_prec1 = checkpoint['best_prec1']
+                model.load_state_dict(checkpoint['state_dict'])
+                optimizer.load_state_dict(checkpoint['optimizer'])
+                print("=> loaded checkpoint '{}' (epoch {})"
+                      .format(args.resume, checkpoint['epoch']))
+            else:
+                print("=> no checkpoint found at '{}'".format(args.resume))
+        resume()
+
+    # Data loading code
+    traindir = os.path.join(args.data, 'train')
+    valdir = os.path.join(args.data, 'val')
+
+    if(args.arch == "inception_v3"):
+        crop_size = 299
+        val_size = 320 # I chose this value arbitrarily, we can adjust.
+    else:
+        crop_size = 224
+        val_size = 256
+
+    train_dataset = datasets.ImageFolder(
+        traindir,
+        transforms.Compose([
+            transforms.RandomResizedCrop(crop_size),
+            transforms.RandomHorizontalFlip(),
+            # transforms.ToTensor(), Too slow
+            # normalize,
+        ]))
+    val_dataset = datasets.ImageFolder(valdir, transforms.Compose([
+            transforms.Resize(val_size),
+            transforms.CenterCrop(crop_size),
+        ]))
+
+    train_sampler = None
+    val_sampler = None
+    if args.distributed:
+        train_sampler = torch.utils.data.distributed.DistributedSampler(train_dataset)
+        val_sampler = torch.utils.data.distributed.DistributedSampler(val_dataset)
+
+    train_loader = torch.utils.data.DataLoader(
+        train_dataset, batch_size=args.batch_size, shuffle=(train_sampler is None),
+        num_workers=args.workers, pin_memory=True, sampler=train_sampler, collate_fn=fast_collate)
+
+    val_loader = torch.utils.data.DataLoader(
+        val_dataset,
+        batch_size=args.batch_size, shuffle=False,
+        num_workers=args.workers, pin_memory=True,
+        sampler=val_sampler,
+        collate_fn=fast_collate)
+
+    if args.evaluate:
+        validate(val_loader, model, criterion)
+        return
+
+    for epoch in range(args.start_epoch, args.epochs):
+        if args.distributed:
+            train_sampler.set_epoch(epoch)
+
+        # train for one epoch
+        train(train_loader, model, criterion, optimizer, epoch)
+        if args.prof:
+            break
+        # evaluate on validation set
+        prec1 = validate(val_loader, model, criterion)
+
+        # remember best prec@1 and save checkpoint
+        if args.local_rank == 0:
+            is_best = prec1 > best_prec1
+            best_prec1 = max(prec1, best_prec1)
+            save_checkpoint({
+                'epoch': epoch + 1,
+                'arch': args.arch,
+                'state_dict': model.state_dict(),
+                'best_prec1': best_prec1,
+                'optimizer' : optimizer.state_dict(),
+            }, is_best)
+
+class data_prefetcher():
+    def __init__(self, loader):
+        self.loader = iter(loader)
+        self.stream = torch.cuda.Stream()
+        self.mean = torch.tensor([0.485 * 255, 0.456 * 255, 0.406 * 255]).cuda().view(1,3,1,1)
+        self.std = torch.tensor([0.229 * 255, 0.224 * 255, 0.225 * 255]).cuda().view(1,3,1,1)
+        # With Amp, it isn't necessary to manually convert data to half.
+        # if args.fp16:
+        #     self.mean = self.mean.half()
+        #     self.std = self.std.half()
+        self.preload()
+
+    def preload(self):
+        try:
+            self.next_input, self.next_target = next(self.loader)
+        except StopIteration:
+            self.next_input = None
+            self.next_target = None
+            return
+        with torch.cuda.stream(self.stream):
+            self.next_input = self.next_input.cuda(non_blocking=True)
+            self.next_target = self.next_target.cuda(non_blocking=True)
+            # With Amp, it isn't necessary to manually convert data to half.
+            # if args.fp16:
+            #     self.next_input = self.next_input.half()
+            # else:
+            self.next_input = self.next_input.float()
+            self.next_input = self.next_input.sub_(self.mean).div_(self.std)
+            
+    def next(self):
+        torch.cuda.current_stream().wait_stream(self.stream)
+        input = self.next_input
+        target = self.next_target
+        self.preload()
+        return input, target
+
+
+def train(train_loader, model, criterion, optimizer, epoch):
+    batch_time = AverageMeter()
+    data_time = AverageMeter()
+    losses = AverageMeter()
+    top1 = AverageMeter()
+    top5 = AverageMeter()
+
+    # switch to train mode
+    model.train()
+    end = time.time()
+
+    run_info_dict = {"Iteration" : [],
+                     "Loss" : [],
+                     "Speed" : []}
+
+    prefetcher = data_prefetcher(train_loader)
+    input, target = prefetcher.next()
+    i = -1
+    while input is not None:
+        i += 1
+
+        # No learning rate warmup for this test, to expose bitwise inaccuracies more quickly
+        # adjust_learning_rate(optimizer, epoch, i, len(train_loader))
+
+        if args.prof:
+            if i > 10:
+                break
+        # measure data loading time
+        data_time.update(time.time() - end)
+
+        # compute output
+        output = model(input)
+        loss = criterion(output, target)
+
+        # measure accuracy and record loss
+        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
+
+        if args.distributed:
+            reduced_loss = reduce_tensor(loss.data)
+            prec1 = reduce_tensor(prec1)
+            prec5 = reduce_tensor(prec5)
+        else:
+            reduced_loss = loss.data
+
+        losses.update(to_python_float(reduced_loss), input.size(0))
+        top1.update(to_python_float(prec1), input.size(0))
+        top5.update(to_python_float(prec5), input.size(0))
+
+        # compute gradient and do SGD step
+        optimizer.zero_grad()
+
+        with amp.scale_loss(loss, optimizer) as scaled_loss:
+            scaled_loss.backward()
+
+        # for param in model.parameters():
+        #     print(param.data.double().sum().item(), param.grad.data.double().sum().item())
+
+        # torch.cuda.synchronize()
+        torch.cuda.nvtx.range_push("step")
+        optimizer.step()
+        torch.cuda.nvtx.range_pop()
+
+        torch.cuda.synchronize()
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+
+        end = time.time()
+
+        # If you decide to refactor this test, like examples/imagenet, to sample the loss every
+        # print_freq iterations, make sure to move this prefetching below the accuracy calculation.
+        input, target = prefetcher.next()
+
+        if i % args.print_freq == 0 and i > 1:
+            if args.local_rank == 0:
+                print('Epoch: [{0}][{1}/{2}]\t'
+                      'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                      'Speed {3:.3f} ({4:.3f})\t'
+                      'Data {data_time.val:.3f} ({data_time.avg:.3f})\t'
+                      'Loss {loss.val:.10f} ({loss.avg:.4f})\t'
+                      'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
+                      'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
+                       epoch, i, len(train_loader),
+                       args.world_size * args.batch_size / batch_time.val,
+                       args.world_size * args.batch_size / batch_time.avg,
+                       batch_time=batch_time,
+                       data_time=data_time, loss=losses, top1=top1, top5=top5))
+            run_info_dict["Iteration"].append(i)
+            run_info_dict["Loss"].append(losses.val)
+            run_info_dict["Speed"].append(args.world_size * args.batch_size / batch_time.val)
+            if len(run_info_dict["Loss"]) == args.prints_to_process:
+                if args.local_rank == 0:
+                    torch.save(run_info_dict,
+                               str(args.has_ext) + "_" + str(args.opt_level) + "_" +
+                               str(args.loss_scale) + "_" + str(args.keep_batchnorm_fp32) + "_" +
+                               str(args.fused_adam))
+                quit()
+
+
+def validate(val_loader, model, criterion):
+    batch_time = AverageMeter()
+    losses = AverageMeter()
+    top1 = AverageMeter()
+    top5 = AverageMeter()
+
+    # switch to evaluate mode
+    model.eval()
+
+    end = time.time()
+
+    prefetcher = data_prefetcher(val_loader)
+    input, target = prefetcher.next()
+    i = -1
+    while input is not None:
+        i += 1
+
+        # compute output
+        with torch.no_grad():
+            output = model(input)
+            loss = criterion(output, target)
+
+        # measure accuracy and record loss
+        prec1, prec5 = accuracy(output.data, target, topk=(1, 5))
+
+        if args.distributed:
+            reduced_loss = reduce_tensor(loss.data)
+            prec1 = reduce_tensor(prec1)
+            prec5 = reduce_tensor(prec5)
+        else:
+            reduced_loss = loss.data
+
+        losses.update(to_python_float(reduced_loss), input.size(0))
+        top1.update(to_python_float(prec1), input.size(0))
+        top5.update(to_python_float(prec5), input.size(0))
+
+        # measure elapsed time
+        batch_time.update(time.time() - end)
+        end = time.time()
+
+        if args.local_rank == 0 and i % args.print_freq == 0:
+            print('Test: [{0}/{1}]\t'
+                  'Time {batch_time.val:.3f} ({batch_time.avg:.3f})\t'
+                  'Speed {2:.3f} ({3:.3f})\t'
+                  'Loss {loss.val:.4f} ({loss.avg:.4f})\t'
+                  'Prec@1 {top1.val:.3f} ({top1.avg:.3f})\t'
+                  'Prec@5 {top5.val:.3f} ({top5.avg:.3f})'.format(
+                   i, len(val_loader),
+                   args.world_size * args.batch_size / batch_time.val,
+                   args.world_size * args.batch_size / batch_time.avg,
+                   batch_time=batch_time, loss=losses,
+                   top1=top1, top5=top5))
+
+        input, target = prefetcher.next()
+
+    print(' * Prec@1 {top1.avg:.3f} Prec@5 {top5.avg:.3f}'
+          .format(top1=top1, top5=top5))
+
+    return top1.avg
+
+
+def save_checkpoint(state, is_best, filename='checkpoint.pth.tar'):
+    torch.save(state, filename)
+    if is_best:
+        shutil.copyfile(filename, 'model_best.pth.tar')
+
+
+class AverageMeter(object):
+    """Computes and stores the average and current value"""
+    def __init__(self):
+        self.reset()
+
+    def reset(self):
+        self.val = 0
+        self.avg = 0
+        self.sum = 0
+        self.count = 0
+
+    def update(self, val, n=1):
+        self.val = val
+        self.sum += val * n
+        self.count += n
+        self.avg = self.sum / self.count
+
+
+def adjust_learning_rate(optimizer, epoch, step, len_epoch):
+    """LR schedule that should yield 76% converged accuracy with batch size 256"""
+    factor = epoch // 30
+
+    if epoch >= 80:
+        factor = factor + 1
+
+    lr = args.lr*(0.1**factor)
+
+    """Warmup"""
+    if epoch < 5:
+        lr = lr*float(1 + step + epoch*len_epoch)/(5.*len_epoch)
+
+    # if(args.local_rank == 0):
+    #     print("epoch = {}, step = {}, lr = {}".format(epoch, step, lr))
+
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = lr
+
+
+def accuracy(output, target, topk=(1,)):
+    """Computes the precision@k for the specified values of k"""
+    maxk = max(topk)
+    batch_size = target.size(0)
+
+    _, pred = output.topk(maxk, 1, True, True)
+    pred = pred.t()
+    correct = pred.eq(target.view(1, -1).expand_as(pred))
+
+    res = []
+    for k in topk:
+        correct_k = correct[:k].view(-1).float().sum(0, keepdim=True)
+        res.append(correct_k.mul_(100.0 / batch_size))
+    return res
+
+
+def reduce_tensor(tensor):
+    rt = tensor.clone()
+    dist.all_reduce(rt, op=dist.reduce_op.SUM)
+    rt /= args.world_size
+    return rt
+
+if __name__ == '__main__':
+    main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/run_test.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/run_test.sh
new file mode 100644
index 0000000000..f4ae06c80e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/common/run_test.sh
@@ -0,0 +1,144 @@
+#!/bin/bash
+
+print_banner() {
+  printf "\n\n\n\e[30m\e[42m$1\e[0m\n\n\n\n"
+}
+
+print_banner "Distributed status:  $1"
+
+echo $2
+DATADIR=$2
+
+if [ -n "$3" ]
+then
+  USE_BASELINE=""
+else
+  USE_BASELINE="--use_baseline"
+fi
+
+if [ "$1" == "single_gpu" ]
+then
+  BASE_CMD="python main_amp.py -a resnet50 --b 128 --workers 4 --deterministic --prints-to-process 5"
+fi
+
+if [ "$1" == "distributed" ]
+then
+  BASE_CMD="python -m torch.distributed.launch --nproc_per_node=2 main_amp.py -a resnet50 --b 128 --workers 4 --deterministic --prints-to-process 5"
+fi
+
+ADAM_ARGS="--opt-level O2 --keep-batchnorm-fp32 False --fused-adam"
+
+keep_batchnorms=(
+""
+"--keep-batchnorm-fp32 True"
+"--keep-batchnorm-fp32 False"
+)
+
+loss_scales=(
+""
+"--loss-scale 1.0"
+"--loss-scale 128.0"
+"--loss-scale dynamic"
+)
+
+opt_levels=(
+"O0"
+"O1"
+"O2"
+"O3"
+)
+
+rm True*
+rm False*
+
+set -e
+
+print_banner "Installing Apex with --cuda_ext and --cpp_ext"
+
+pushd ../../..
+pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .
+popd
+
+for opt_level in "${opt_levels[@]}"
+do
+  for loss_scale in "${loss_scales[@]}"
+  do
+    for keep_batchnorm in "${keep_batchnorms[@]}"
+    do
+      if [ "$opt_level" == "O1" ] && [ -n "${keep_batchnorm}" ]
+      then
+        print_banner "Skipping ${opt_level} ${loss_scale} ${keep_batchnorm}"
+        continue
+      fi
+      print_banner "${BASE_CMD} --opt-level ${opt_level} ${loss_scale} ${keep_batchnorm} --has-ext $DATADIR"
+      set -x
+      ${BASE_CMD} --opt-level ${opt_level} ${loss_scale} ${keep_batchnorm} --has-ext $DATADIR
+      set +x
+    done
+  done
+done
+
+# Handle FusedAdam separately due to limited support.
+# FusedAdam will not be tested for bitwise accuracy against the Python implementation.
+# The L0 tests already do so.  These tests are here to ensure that it actually runs,
+# and get an idea of performance.
+for loss_scale in "${loss_scales[@]}"
+do
+  print_banner "${BASE_CMD} ${ADAM_ARGS} ${loss_scale} --has-ext $DATADIR"
+  set -x
+  ${BASE_CMD} ${ADAM_ARGS} ${loss_scale} --has-ext $DATADIR
+  set +x
+done
+
+print_banner "Reinstalling apex without extensions"
+
+pushd ../../..
+pip install -v --no-cache-dir .
+popd
+
+for opt_level in "${opt_levels[@]}"
+do
+  for loss_scale in "${loss_scales[@]}"
+  do
+    for keep_batchnorm in "${keep_batchnorms[@]}"
+    do
+      if [ "$opt_level" == "O1" ] && [ -n "${keep_batchnorm}" ]
+      then
+        print_banner "Skipping ${opt_level} ${loss_scale} ${keep_batchnorm}"
+        continue
+      fi
+      print_banner "${BASE_CMD} --opt-level ${opt_level} ${loss_scale} ${keep_batchnorm} $DATADIR"
+      set -x
+      ${BASE_CMD} --opt-level ${opt_level} ${loss_scale} ${keep_batchnorm} $DATADIR
+      set +x
+    done
+  done
+done
+
+print_banner "Checking for bitwise accuracy between Python-only and cpp/cuda extension installs"
+
+for opt_level in "${opt_levels[@]}"
+do
+  for loss_scale in "${loss_scales[@]}"
+  do
+    for keep_batchnorm in "${keep_batchnorms[@]}"
+    do
+      echo ""
+      if [ "$opt_level" == "O1" ] && [ -n "${keep_batchnorm}" ]
+      then
+        echo "Skipping ${opt_level} ${loss_scale} ${keep_batchnorm}"
+        continue
+      fi
+      echo "${BASE_CMD} --opt-level ${opt_level} ${loss_scale} ${keep_batchnorm} [--has-ext] $DATADIR"
+      set -x
+      python compare.py --opt-level ${opt_level} ${loss_scale} ${keep_batchnorm} --use_baseline
+      set +x
+    done
+  done
+done
+
+print_banner "Reinstalling Apex with --cuda_ext and --cpp_ext"
+
+pushd ../../..
+pip install -v --no-cache-dir --global-option="--cpp_ext" --global-option="--cuda_ext" .
+popd
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/cross_product/run.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/cross_product/run.sh
new file mode 100644
index 0000000000..7ccf9ec438
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/cross_product/run.sh
@@ -0,0 +1,6 @@
+#!/bin/bash
+
+# DATADIR="/home/mcarilli/Desktop/pt18data/apex_stale/examples/imagenet/bare_metal_train_val/"
+# DATADIR="/opt/home/apex/examples/imagenet/"
+cp ../common/* .
+bash run_test.sh single_gpu $1
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/cross_product_distributed/run.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/cross_product_distributed/run.sh
new file mode 100644
index 0000000000..917ec11e83
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/L1/cross_product_distributed/run.sh
@@ -0,0 +1,4 @@
+#!/bin/bash
+
+cp ../common/* .
+bash run_test.sh distributed $1
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/DDP/ddp_race_condition_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/DDP/ddp_race_condition_test.py
new file mode 100644
index 0000000000..761a33595c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/DDP/ddp_race_condition_test.py
@@ -0,0 +1,69 @@
+import torch
+import torch.distributed as dist
+from torch.nn import Parameter
+from torch.nn import Module
+from apex.parallel import DistributedDataParallel as DDP
+import argparse
+import os
+
+
+parser = argparse.ArgumentParser(description='allreduce hook example')
+parser.add_argument("--local_rank", default=0, type=int)
+args = parser.parse_args()
+
+args.distributed = False
+if 'WORLD_SIZE' in os.environ:
+    args.distributed = int(os.environ['WORLD_SIZE']) > 1
+
+if args.distributed:
+    args.gpu = args.local_rank % torch.cuda.device_count()
+    torch.cuda.set_device(args.gpu)
+    torch.distributed.init_process_group(backend='nccl',
+                                         init_method='env://')
+    args.world_size = torch.distributed.get_world_size()
+
+torch.set_printoptions(precision=10)
+torch.manual_seed(args.local_rank)
+
+class Model(Module):
+    def __init__(self):
+        super(Model, self).__init__()
+        self.a = Parameter(torch.cuda.FloatTensor(4096*4096).fill_(1.0))
+        self.b = Parameter(torch.cuda.FloatTensor(4096*4096).fill_(2.0))
+    def forward(self, input):
+        return (input*self.a)*self.b
+
+model = Model()
+# model = DDP(model, message_size=1, gradient_predivide_factor=8.0)
+# model = DDP(model, delay_allreduce=True)
+# model = DDP(model, message_size=1, allreduce_trigger_params=[model.b])
+model = DDP(model, message_size=1, allreduce_trigger_params=[model.b], num_allreduce_streams=3)
+
+x = torch.cuda.FloatTensor(4096*4096)
+
+passed = True
+torch.cuda.cudart().cudaProfilerStart()
+for i in range(10):
+    x.fill_(i + args.local_rank) # fill x with new values every iteration for sanity
+    model.zero_grad()
+    out = model(x)
+    loss = out.sum()
+    # torch.cuda.nvtx.range_push("backward")
+    loss.backward()
+    # torch.cuda.nvtx.range_pop()
+    
+    # torch.cuda.nvtx.range_push("synchronize() + info")
+    # torch.cuda.synchronize()
+    print("i = {}".format(i))
+    def info(name, param, val):
+        expected = val*4096*4096*(2.*i+1)/2.
+        actual = param.grad.data.sum().item()
+        print(name+": grad.data_ptr() = {}, expected sum {}, got {}".format(
+              param.grad.data_ptr(), expected, actual))
+        return (expected == actual)
+    if not info("model.a", model.module.a, 2.):  passed = False
+    if not info("model.b", model.module.b, 1.):  passed = False
+    # torch.cuda.nvtx.range_pop()
+torch.cuda.cudart().cudaProfilerStop()
+
+print("passed = ", passed)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/DDP/run_race_test.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/DDP/run_race_test.sh
new file mode 100644
index 0000000000..2c2bd266cc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/DDP/run_race_test.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+CUDA_VISIBLE_DEVICES=0,1 python -m torch.distributed.launch --nproc_per_node=2 ddp_race_condition_test.py
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/amp_master_params.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/amp_master_params.py
new file mode 100644
index 0000000000..4af5092f75
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/amp_master_params.py
@@ -0,0 +1,70 @@
+import torch
+import argparse
+import os
+from apex import amp
+# FOR DISTRIBUTED: (can also use torch.nn.parallel.DistributedDataParallel instead)
+from apex.parallel import DistributedDataParallel
+
+parser = argparse.ArgumentParser()
+# FOR DISTRIBUTED:  Parse for the local_rank argument, which will be supplied
+# automatically by torch.distributed.launch.
+parser.add_argument("--local_rank", default=0, type=int)
+args = parser.parse_args()
+
+# FOR DISTRIBUTED:  If we are running under torch.distributed.launch,
+# the 'WORLD_SIZE' environment variable will also be set automatically.
+args.distributed = False
+if 'WORLD_SIZE' in os.environ:
+    args.distributed = int(os.environ['WORLD_SIZE']) > 1
+
+if args.distributed:
+    # FOR DISTRIBUTED:  Set the device according to local_rank.
+    torch.cuda.set_device(args.local_rank)
+
+    # FOR DISTRIBUTED:  Initialize the backend.  torch.distributed.launch will provide
+    # environment variables, and requires that you use init_method=`env://`.
+    torch.distributed.init_process_group(backend='nccl',
+                                         init_method='env://')
+
+    torch.manual_seed(torch.distributed.get_rank())
+
+torch.backends.cudnn.benchmark = True
+
+N, D_in, D_out = 64, 1024, 16
+
+# Each process receives its own batch of "fake input data" and "fake target data."
+# The "training loop" in each process just uses this fake batch over and over.
+# https://github.com/NVIDIA/apex/tree/master/examples/imagenet provides a more realistic
+# example of distributed data sampling for both training and validation.
+x = torch.randn(N, D_in, device='cuda')
+y = torch.randn(N, D_out, device='cuda')
+
+model = torch.nn.Linear(D_in, D_out).cuda()
+optimizer = torch.optim.SGD(model.parameters(), lr=1e-3)
+
+model, optimizer = amp.initialize(model, optimizer, opt_level="O2")
+
+if args.distributed:
+    # FOR DISTRIBUTED:  After amp.initialize, wrap the model with
+    # apex.parallel.DistributedDataParallel.
+    model = DistributedDataParallel(model)
+    # torch.nn.parallel.DistributedDataParallel is also fine, with some added args:
+    # model = torch.nn.parallel.DistributedDataParallel(model,
+    #                                                   device_ids=[args.local_rank],
+    #                                                   output_device=args.local_rank)
+
+loss_fn = torch.nn.MSELoss()
+
+for t in range(500):
+    optimizer.zero_grad()
+    y_pred = model(x)
+    loss = loss_fn(y_pred, y)
+    with amp.scale_loss(loss, optimizer) as scaled_loss:
+        scaled_loss.backward()
+    optimizer.step()
+
+if args.local_rank == 0:
+    print("final loss = ", loss)
+
+torch.save(list(model.parameters()), "rank{}model.pth".format(torch.distributed.get_rank()))
+torch.save(list(amp.master_params(optimizer)), "rank{}master.pth".format(torch.distributed.get_rank()))
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/compare.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/compare.py
new file mode 100644
index 0000000000..e5cbf20c1d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/compare.py
@@ -0,0 +1,28 @@
+import torch
+
+model_params_rank0 = torch.load("rank0model.pth",
+                           map_location = lambda storage, loc: storage.cuda(0))
+model_params_rank1 = torch.load("rank1model.pth",
+                                 map_location = lambda storage, loc: storage.cuda(0))
+master_params_rank0 = torch.load("rank0master.pth",
+                                 map_location = lambda storage, loc: storage.cuda(0))
+master_params_rank1 = torch.load("rank1master.pth",
+                                 map_location = lambda storage, loc: storage.cuda(0))
+
+for model_rank0, model_rank1, master_rank0, master_rank1 in zip(
+        model_params_rank0,
+        model_params_rank1,
+        master_params_rank0,
+        master_params_rank1):
+    assert torch.allclose(model_rank0, model_rank1), "Model param mismatch"
+    assert torch.allclose(master_rank0, master_rank1), "Master param mismatch"
+    # Some debugging/investigation assistance code:
+    # maxval, maxind = torch.max(((torch.abs(model_rank0).float())/torch.abs(master_rank0)).view(-1), 0)
+    # offending_val_half = model_rank0.view(-1)[maxind.item()]
+    # offending_val_float = master_rank0.view(-1)[maxind.item()]
+    # print(maxval.item(), maxind.item(), offending_val_half.item(), offending_val_float.item(),
+    #       offending_val_float.half().item())
+    # rtol needs to be > 2^-11 because of denormals...
+    assert torch.allclose(model_rank0, master_rank0.half(), rtol=.005), "Model-master mismatch"
+
+print("OK:  Model and master params match across ranks.")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/run.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/run.sh
new file mode 100644
index 0000000000..8599dbbbc9
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/amp_master_params/run.sh
@@ -0,0 +1,4 @@
+#!/bin/bash
+python -m torch.distributed.launch --nproc_per_node=2 amp_master_params.py
+
+python compare.py
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/python_single_gpu_unit_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/python_single_gpu_unit_test.py
new file mode 100644
index 0000000000..80c26825ba
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/python_single_gpu_unit_test.py
@@ -0,0 +1,111 @@
+import torch
+import numpy as np
+import apex
+
+def compare(desc, inp1, inp2, error):
+    a = inp1.clone().detach().cpu().numpy()
+    b = inp2.clone().detach().cpu().numpy()
+    close = np.allclose(a,b, error, error)
+    if not close:
+        print(desc, close)
+        z = a - b
+        index = (np.abs(z) >= error + error * np.abs(b)).nonzero()
+        print("dif    : ", z[index])
+        print("inp1   : ", a[index])
+        print("inp2   : ", b[index])
+    return close
+
+feature_size = 10
+space_size = 16
+batch_size = 5
+
+
+error = 1e-5
+
+np.random.seed(1)
+dtype = np.float32
+inp = (np.random.randn(batch_size, feature_size, space_size, space_size)).astype(dtype)
+grad = (np.random.randn(batch_size, feature_size, space_size, space_size)).astype(dtype)
+weight = (np.random.randn(feature_size)).astype(dtype)
+bias = (np.random.randn(feature_size)).astype(dtype)
+
+type_tensor = torch.cuda.FloatTensor
+ref_tensor = torch.cuda.DoubleTensor
+
+inp_t = type_tensor(inp)
+weight_t = type_tensor(weight)
+bias_t = type_tensor(bias)
+
+inp_r = ref_tensor(inp.transpose(1, 0, 2, 3).reshape(feature_size, -1))
+inp2_r = ref_tensor(inp)
+weight_r = ref_tensor(weight).view(-1, 1, 1)
+bias_r = ref_tensor(bias).view(-1, 1, 1)
+
+grad_output_t = type_tensor(grad)
+
+m = inp_r.mean(1)
+b_v = inp_r.var(1, unbiased=False)
+unb_v = inp_r.var(1, unbiased=True)
+
+eps = 1e-5
+
+bn = torch.nn.BatchNorm2d(feature_size).cuda()
+bn.momentum = 1.0
+bn.weight.data = weight_t.clone()
+bn.bias.data = bias_t.clone()
+inp_bn = inp_t.clone().requires_grad_()
+grad_bn = grad_output_t.clone().detach()
+out_bn = bn(inp_bn)
+out_bn.backward(grad_bn)
+
+from apex.parallel.sync_batchnorm import SyncBatchNorm
+
+sbn = SyncBatchNorm(feature_size).cuda()
+sbn.momentum = 1.0
+sbn.weight.data = weight_t.clone()
+sbn.bias.data = bias_t.clone()
+inp_sbn = inp_t.clone().requires_grad_()
+grad_sbn = grad_output_t.clone().detach()
+out_sbn = sbn(inp_sbn)
+out_sbn.backward(grad_sbn)
+
+sbn_result = True
+sbn_result_c_last = True
+bn_result = True
+
+out_r = weight_r * (inp2_r - m.view(-1, 1, 1)) * torch.rsqrt(b_v.view(-1,1,1) + eps) + bias_r
+
+compare("comparing bn output: ", out_bn, out_r, error)
+
+grad_output_t = type_tensor(grad)
+
+grad_output_r = ref_tensor(grad.transpose(1, 0, 2, 3).reshape(feature_size, -1))
+grad_output2_r = ref_tensor(grad)
+
+grad_bias_r = grad_output_r.sum(1)
+grad_weight_r = ((inp2_r - m.view(-1, 1, 1)) * torch.rsqrt(b_v.view(-1,1,1) + eps) * grad_output2_r).transpose(1,0).contiguous().view(feature_size, -1).sum(1)
+
+mean_dy_r = grad_output_r.mean(1)
+mean_dy_xmu_r = ((inp2_r - m.view(-1, 1, 1)) * grad_output2_r).transpose(1,0).contiguous().view(feature_size, -1).mean(1)
+
+grad_input_r = (grad_output2_r - mean_dy_r.view(-1, 1, 1) - (inp2_r - m.view(-1, 1, 1)) / (b_v.view(-1,1,1) + eps) * mean_dy_xmu_r.view(-1, 1, 1) ) * torch.rsqrt(b_v.view(-1,1,1) + eps) * weight_r.view(-1,1,1)
+
+compare("comparing bn input grad: ", inp_bn.grad, grad_input_r, error)
+sbn_result = compare("comparing sbn input grad: ", inp_sbn.grad, grad_input_r, error) and sbn_result
+
+compare("comparing bn/sbn output: ", out_bn, out_sbn, error)
+sbn_result = compare("comparing running_mean: ", bn.running_mean.data, sbn.running_mean.data, error) and sbn_result
+sbn_result = compare("comparing running_variance: ", bn.running_var.data, sbn.running_var.data, error) and sbn_result
+compare("comparing grad_input: ", inp_bn.grad, inp_sbn.grad, error)
+compare("comparing grad_bias: ", bn.bias.grad, sbn.bias.grad, error)
+compare("comparing grad_bias bn to ref: ", bn.bias.grad, grad_bias_r, error)
+sbn_result = compare("comparing grad_bias sbn to ref: ", sbn.bias.grad, grad_bias_r, error) and sbn_result
+compare("comparing grad_weight: ", bn.weight.grad, sbn.weight.grad, error)
+compare("comparing grad_weight bn to ref: ", bn.weight.grad, grad_weight_r, error)
+sbn_result = compare("comparing grad_weight sbn to ref: ", sbn.weight.grad, grad_weight_r, error) and sbn_result
+
+if sbn_result:
+    print("====SBN single gpu passed tests")
+else:
+    print("*SBN single gpu failed*")
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/single_gpu_unit_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/single_gpu_unit_test.py
new file mode 100644
index 0000000000..6fcbd0d147
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/single_gpu_unit_test.py
@@ -0,0 +1,159 @@
+import torch
+import numpy as np
+import apex
+if True:
+    print("using setup tools")
+    import syncbn
+else:
+    print("using jit")
+    from torch.utils.cpp_extension import load
+    syncbn = load(name='syncbn', sources=['../../csrc/syncbn.cpp', '../../csrc/welford.cu'])
+
+def compare(desc, inp1, inp2, error):
+    a = inp1.clone().detach().cpu().numpy()
+    b = inp2.clone().detach().cpu().numpy()
+    close = np.allclose(a,b, error, error)
+    if not close:
+        print(desc, close)
+        z = a - b
+        index = (np.abs(z) >= error + error * np.abs(b)).nonzero()
+        print("dif    : ", z[index])
+        print("inp1   : ", a[index])
+        print("inp2   : ", b[index])
+    return close
+
+feature_size = 10
+space_size = 16
+batch_size = 5
+
+
+error = 1e-5
+
+np.random.seed(1)
+dtype = np.float32
+inp = (np.random.randn(batch_size, feature_size, space_size, space_size)).astype(dtype)
+grad = (np.random.randn(batch_size, feature_size, space_size, space_size)).astype(dtype)
+weight = (np.random.randn(feature_size)).astype(dtype)
+bias = (np.random.randn(feature_size)).astype(dtype)
+count = torch.cuda.IntTensor([batch_size*space_size**2])
+
+type_tensor = torch.cuda.FloatTensor
+ref_tensor = torch.cuda.DoubleTensor
+
+inp_t = type_tensor(inp)
+weight_t = type_tensor(weight)
+bias_t = type_tensor(bias)
+
+inp_r = ref_tensor(inp.transpose(1, 0, 2, 3).reshape(feature_size, -1))
+inp2_r = ref_tensor(inp)
+weight_r = ref_tensor(weight).view(-1, 1, 1)
+bias_r = ref_tensor(bias).view(-1, 1, 1)
+
+grad_output_t = type_tensor(grad)
+
+m = inp_r.mean(1)
+b_v = inp_r.var(1, unbiased=False)
+unb_v = inp_r.var(1, unbiased=True)
+
+eps = 1e-5
+
+#mean, var, var_biased = syncbn.welford_mean_var(inp_t)
+mean, var_biased = syncbn.welford_mean_var(inp_t)
+inv_std = 1.0 / torch.sqrt(var_biased + eps)
+
+bn = torch.nn.BatchNorm2d(feature_size).cuda()
+bn.momentum = 1.0
+bn.weight.data = weight_t.clone()
+bn.bias.data = bias_t.clone()
+inp_bn = inp_t.clone().requires_grad_()
+grad_bn = grad_output_t.clone().detach()
+out_bn = bn(inp_bn)
+out_bn.backward(grad_bn)
+
+sbn = apex.parallel.SyncBatchNorm(feature_size).cuda()
+sbn.momentum = 1.0
+sbn.weight.data = weight_t.clone()
+sbn.bias.data = bias_t.clone()
+inp_sbn = inp_t.clone().requires_grad_()
+grad_sbn = grad_output_t.clone().detach()
+out_sbn = sbn(inp_sbn)
+out_sbn.backward(grad_sbn)
+
+sbn_c_last = apex.parallel.SyncBatchNorm(feature_size, channel_last=True).cuda()
+sbn_c_last.momentum = 1.0
+sbn_c_last.weight.data = weight_t.clone()
+sbn_c_last.bias.data = bias_t.clone()
+inp_sbn_c_last = inp_t.clone().transpose(-1, 1).contiguous().requires_grad_()
+grad_sbn_c_last = grad_output_t.clone().transpose(-1, 1).contiguous().detach()
+out_sbn_c_last = sbn_c_last(inp_sbn_c_last)
+out_sbn_c_last.backward(grad_sbn_c_last)
+
+sbn_result = True
+sbn_result_c_last = True
+bn_result = True
+
+sbn_result = compare("comparing mean: ", mean, m, error) and sbn_result
+#sbn_result = compare("comparing variance: ", var, unb_v, error) and sbn_result
+sbn_result = compare("comparing biased variance: ", var_biased, b_v, error) and sbn_result
+
+
+out = syncbn.batchnorm_forward(inp_t, mean, inv_std, weight_t, bias_t)
+out_r = weight_r * (inp2_r - m.view(-1, 1, 1)) * torch.rsqrt(b_v.view(-1,1,1) + eps) + bias_r
+
+sbn_result = compare("comparing output: ", out, out_r, error) and sbn_result
+compare("comparing bn output: ", out_bn, out_r, error)
+
+grad_output_t = type_tensor(grad)
+
+grad_output_r = ref_tensor(grad.transpose(1, 0, 2, 3).reshape(feature_size, -1))
+grad_output2_r = ref_tensor(grad)
+
+grad_bias_r = grad_output_r.sum(1)
+grad_weight_r = ((inp2_r - m.view(-1, 1, 1)) * torch.rsqrt(b_v.view(-1,1,1) + eps) * grad_output2_r).transpose(1,0).contiguous().view(feature_size, -1).sum(1)
+
+sum_dy_r = grad_output_r.sum(1)
+mean_dy_r = grad_output_r.mean(1)
+sum_dy_xmu_r = ((inp2_r - m.view(-1, 1, 1)) * grad_output2_r).transpose(1,0).contiguous().view(feature_size, -1).sum(1)
+mean_dy_xmu_r = ((inp2_r - m.view(-1, 1, 1)) * grad_output2_r).transpose(1,0).contiguous().view(feature_size, -1).mean(1)
+
+grad_input_r = (grad_output2_r - mean_dy_r.view(-1, 1, 1) - (inp2_r - m.view(-1, 1, 1)) / (b_v.view(-1,1,1) + eps) * mean_dy_xmu_r.view(-1, 1, 1) ) * torch.rsqrt(b_v.view(-1,1,1) + eps) * weight_r.view(-1,1,1)
+
+sum_dy, sum_dy_xmu, grad_weight, grad_bias = syncbn.reduce_bn(grad_output_t, inp_t, mean, inv_std, weight_t)
+grad_input = syncbn.batchnorm_backward(grad_output_t, inp_t, mean, inv_std, weight_t, sum_dy, sum_dy_xmu, count)
+sbn_result = compare("comparing bias grad: ", grad_bias, grad_bias_r, error) and sbn_result
+sbn_result = compare("comparing weight grad: ", grad_weight, grad_weight_r, error) and sbn_result
+sbn_result = compare("comparing sum_dy grad: ", sum_dy, sum_dy_r, error) and sbn_result
+sbn_result = compare("comparing sum_dy_xmu grad: ", sum_dy_xmu, sum_dy_xmu_r, error) and sbn_result
+sbn_result = compare("comparing input grad: ", grad_input, grad_input_r, error) and sbn_result
+compare("comparing bn input grad: ", inp_bn.grad, grad_input_r, error)
+sbn_result = compare("comparing sbn input grad: ", inp_sbn.grad, grad_input_r, error) and sbn_result
+
+compare("comparing bn/sbn output: ", out_bn, out_sbn, error)
+sbn_result = compare("comparing running_mean: ", bn.running_mean.data, sbn.running_mean.data, error) and sbn_result
+sbn_result = compare("comparing running_variance: ", bn.running_var.data, sbn.running_var.data, error) and sbn_result
+compare("comparing grad_input: ", inp_bn.grad, inp_sbn.grad, error)
+compare("comparing grad_bias: ", bn.bias.grad, sbn.bias.grad, error)
+compare("comparing grad_bias bn to ref: ", bn.bias.grad, grad_bias_r, error)
+sbn_result = compare("comparing grad_bias sbn to ref: ", sbn.bias.grad, grad_bias_r, error) and sbn_result
+compare("comparing grad_weight: ", bn.weight.grad, sbn.weight.grad, error)
+compare("comparing grad_weight bn to ref: ", bn.weight.grad, grad_weight_r, error)
+sbn_result = compare("comparing grad_weight sbn to ref: ", sbn.weight.grad, grad_weight_r, error) and sbn_result
+
+compare("comparing channel last bn/sbn output: ", out_bn, out_sbn_c_last.transpose(-1, 1).contiguous(), error)
+sbn_result_c_last = compare("comparing channel last running_mean: ", bn.running_mean.data, sbn_c_last.running_mean.data, error) and sbn_result_c_last
+sbn_result_c_last = compare("comparing channel last running_variance: ", bn.running_var.data, sbn_c_last.running_var.data, error) and sbn_result_c_last
+compare("comparing channel last grad_input: ", inp_bn.grad, inp_sbn_c_last.grad.transpose(-1, 1).contiguous(), error)
+compare("comparing channel last grad_bias: ", bn.bias.grad, sbn_c_last.bias.grad, error)
+sbn_result_c_last = compare("comparing channel last grad_bias sbn to ref: ", sbn_c_last.bias.grad, grad_bias_r, error) and sbn_result_c_last
+compare("comparing channel last grad_weight: ", bn.weight.grad, sbn_c_last.weight.grad, error)
+sbn_result_c_last = compare("comparing channel last grad_weight sbn to ref: ", sbn_c_last.weight.grad, grad_weight_r, error) and sbn_result_c_last
+
+if sbn_result:
+    print("====SBN single gpu passed tests")
+else:
+    print("*SBN single gpu failed*")
+
+if sbn_result_c_last:
+    print("====SBN channel last single gpu passed tests")
+else:
+    print("*SBN channel last single gpu failed*")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/test_batchnorm1d.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/test_batchnorm1d.py
new file mode 100644
index 0000000000..f35ac47397
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/test_batchnorm1d.py
@@ -0,0 +1,18 @@
+import torch
+import apex
+
+model = apex.parallel.SyncBatchNorm(4).cuda()
+model.weight.data.uniform_()
+model.bias.data.uniform_()
+data = torch.rand((8,4)).cuda()
+
+model_ref = torch.nn.BatchNorm1d(4).cuda()
+model_ref.load_state_dict(model.state_dict())
+data_ref = data.clone()
+
+output = model(data)
+output_ref = model_ref(data_ref)
+
+assert(output.allclose(output_ref))
+assert(model.running_mean.allclose(model_ref.running_mean))
+assert(model.running_var.allclose(model_ref.running_var))
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/test_groups.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/test_groups.py
new file mode 100644
index 0000000000..d028cc397f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/test_groups.py
@@ -0,0 +1,185 @@
+import torch
+import numpy as np
+import apex
+import syncbn
+import os
+import argparse
+import torch.optim as optim
+
+def compare(desc, inp1, inp2, error):
+    a = inp1.clone().detach().cpu().numpy()
+    b = inp2.clone().detach().cpu().numpy()
+    close = np.allclose(a,b, error, error)
+    if not close:
+        print(desc, close)
+        z = a - b
+        index = (np.abs(z) >= error + error * np.abs(b)).nonzero()
+        print("dif    : ", z[index])
+        print("inp1   : ", a[index])
+        print("inp2   : ", b[index])
+    return close
+
+feature_size = 10
+space_size = 40
+batch_size = 32
+
+
+from apex.parallel import DistributedDataParallel as DDP
+parser = argparse.ArgumentParser()
+parser.add_argument("--local_rank", default=0, type=int)
+parser.add_argument("--fp16", action='store_true', default=False)
+parser.add_argument("--fp64", action='store_true', default=False)
+parser.add_argument("--group_size", default=0, type=int)
+args = parser.parse_args()
+
+try:
+    args.world_size = int(os.environ['WORLD_SIZE'])
+except:
+    print("This is a multi-gpu test. To run it please use 'python -m torch.distributed.launch --nproc_per_node=<num gpus> test_groups.py <more options>'")
+    exit(1)
+
+torch.cuda.set_device(args.local_rank)
+torch.distributed.init_process_group(backend='nccl', init_method='env://')
+
+start = (args.local_rank%args.group_size) * batch_size//args.group_size
+finish = (args.local_rank%args.group_size + 1) * batch_size//args.group_size
+
+error = 1e-5
+dtype = np.float32
+if args.fp16:
+    error = 1e-3
+    dtype = np.float16
+elif args.fp64:
+    error = 1e-8
+    dtype = np.float64
+
+
+np.random.seed(18 + args.local_rank//args.group_size)
+
+inp = np.random.randn(batch_size, feature_size, space_size, space_size).astype(dtype)
+grad = np.random.randn(batch_size, feature_size, space_size, space_size).astype(dtype)
+weight = np.random.randn(feature_size).astype(dtype)
+bias = np.random.randn(feature_size).astype(dtype)
+
+
+type_tensor = torch.cuda.FloatTensor
+if args.fp16:
+    type_tensor = torch.cuda.HalfTensor
+if args.fp64:
+    type_tensor = torch.cuda.DoubleTensor
+
+ref_tensor = torch.cuda.DoubleTensor
+
+inp_t = type_tensor(inp)
+weight_t = type_tensor(weight)
+bias_t = type_tensor(bias)
+
+inp_r = ref_tensor(inp.transpose(1, 0, 2, 3).reshape(feature_size, -1))
+inp2_r = ref_tensor(inp)
+weight_r = ref_tensor(weight).view(-1, 1, 1)
+bias_r = ref_tensor(bias).view(-1, 1, 1)
+
+grad_output_t = type_tensor(grad)
+
+m = inp_r.mean(1)
+b_v = inp_r.var(1, unbiased=False)
+unb_v = inp_r.var(1, unbiased=True)
+
+eps = 1e-5
+
+mean, var_biased = syncbn.welford_mean_var(inp_t)
+inv_std = 1.0 / torch.sqrt(var_biased + eps)
+
+bn = torch.nn.BatchNorm2d(feature_size).cuda()
+bn.momentum = 1.0
+bn.weight.data = weight_t.clone()
+bn.bias.data = bias_t.clone()
+if args.fp16:
+    bn.half()
+if args.fp64:
+    bn.double()
+bn = DDP(bn)
+inp_bn = inp_t.clone().requires_grad_()
+grad_bn = grad_output_t.clone().detach()
+out_bn = bn(inp_bn)
+out_bn.backward(grad_bn)
+# compensating the averaging over processes done by DDP
+# in order to produce mathematically equivalent result
+# https://github.com/NVIDIA/apex/issues/134#issuecomment-458307368
+for param in bn.parameters():
+    param.grad = param.grad / args.group_size
+bn_opt = optim.SGD(bn.parameters(), lr=1.0)
+
+sbn = apex.parallel.SyncBatchNorm(feature_size, process_group=apex.parallel.create_syncbn_process_group(args.group_size)).cuda()
+sbn.momentum = 1.0
+sbn.weight.data = weight_t.clone()
+sbn.bias.data = bias_t.clone()
+if args.fp16:
+    sbn.half()
+if args.fp64:
+    sbn.double()
+sbn = DDP(sbn)
+sbn_opt = optim.SGD(sbn.parameters(), lr=1.0)
+inp_sbn = inp_t.clone().requires_grad_()
+grad_sbn = grad_output_t.clone().detach()
+out_sbn = sbn(inp_sbn[start:finish])
+out_sbn.backward(grad_sbn[start:finish])
+
+sbn_result = True
+bn_result = True
+
+if args.local_rank == 0:
+    sbn_result = compare("comparing mean: ", mean, m, error) and sbn_result
+    sbn_result = compare("comparing biased variance: ", var_biased, b_v, error) and sbn_result
+
+out = syncbn.batchnorm_forward(inp_t, mean, inv_std, weight_t, bias_t)
+out_r = weight_r * (inp2_r - m.view(-1, 1, 1)) * torch.rsqrt(b_v.view(-1,1,1) + eps) + bias_r
+
+if args.local_rank == 0:
+    sbn_result = compare("comparing output: ", out, out_r, error) and sbn_result
+    compare("comparing bn output: ", out_bn, out_r, error)
+
+grad_output_t = type_tensor(grad)
+
+grad_output_r = ref_tensor(grad.transpose(1, 0, 2, 3).reshape(feature_size, -1))
+grad_output2_r = ref_tensor(grad)
+
+grad_bias_r = grad_output_r.sum(1)
+grad_weight_r = ((inp2_r - m.view(-1, 1, 1)) * torch.rsqrt(b_v.view(-1,1,1) + eps) * grad_output2_r).transpose(1,0).contiguous().view(feature_size, -1).sum(1)
+
+mean_dy_r = grad_output_r.mean(1)
+mean_dy_xmu_r = ((inp2_r - m.view(-1, 1, 1)) * grad_output2_r).transpose(1,0).contiguous().view(feature_size, -1).mean(1)
+
+grad_input_r = (grad_output2_r - mean_dy_r.view(-1, 1, 1) - (inp2_r - m.view(-1, 1, 1)) / (b_v.view(-1,1,1) + eps) * mean_dy_xmu_r.view(-1, 1, 1) ) * torch.rsqrt(b_v.view(-1,1,1) + eps) * weight_r.view(-1,1,1)
+
+mean_dy, mean_dy_xmu, grad_weight, grad_bias = syncbn.reduce_bn(grad_output_t, inp_t, mean, inv_std, weight_t)
+grad_input = syncbn.batchnorm_backward(grad_output_t, inp_t, mean, inv_std, weight_t, mean_dy, mean_dy_xmu)
+
+if args.local_rank == 0:
+    sbn_result = compare("comparing bias grad: ", grad_bias, grad_bias_r, error) and sbn_result
+    sbn_result = compare("comparing weight grad: ", grad_weight, grad_weight_r, error) and sbn_result
+    sbn_result = compare("comparing mean_dy grad: ", mean_dy, mean_dy_r, error) and sbn_result
+    sbn_result = compare("comparing mean_dy_xmu grad: ", mean_dy_xmu, mean_dy_xmu_r, error) and sbn_result
+    sbn_result = compare("comparing input grad: ", grad_input, grad_input_r, error) and sbn_result
+    compare("comparing bn input grad: ", inp_bn.grad, grad_input_r, error)
+
+if args.local_rank == 0:
+    sbn_result = compare("comparing running_mean: ", bn.module.running_mean.data, sbn.module.running_mean.data, error) and sbn_result
+    sbn_result = compare("comparing running_variance: ", bn.module.running_var.data, sbn.module.running_var.data, error) and sbn_result
+
+# execute by both
+compare("comparing layers output: ", out_bn[start:finish], out_sbn, error) and sbn_result
+compare("comparing layers grad_input: ", inp_bn.grad[start:finish], inp_sbn.grad[start:finish], error) and sbn_result
+
+bn_opt.step()
+sbn_opt.step()
+
+if args.local_rank == 0:
+    compare("comparing bn vs sbn bias: ", bn.module.bias, sbn.module.bias, error)
+    compare("comparing bn vs sbn weight: ", bn.module.weight, sbn.module.weight, error)
+
+
+if sbn_result:
+    print("====SBN group test passed")
+else:
+    print("*SBN group test failed*")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/two_gpu_test_different_batch_size.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/two_gpu_test_different_batch_size.py
new file mode 100644
index 0000000000..a9e8cb641a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/two_gpu_test_different_batch_size.py
@@ -0,0 +1,158 @@
+import torch
+import torch.nn as nn
+from torch.nn.parallel import DistributedDataParallel as DDP
+from apex.parallel import SyncBatchNorm as ApexSyncBatchNorm
+
+import argparse
+import os
+import numpy as np
+
+var_batch = 16
+
+def compare(desc, inp1, inp2, error= 1e-5):
+    a = inp1.clone().detach().cpu().numpy()
+    b = inp2.clone().detach().cpu().numpy()
+    close = np.allclose(a,b, error, error)
+    if not close:
+        print(desc, close)
+        z = a - b
+        index = (np.abs(z) >= error + error * np.abs(b)).nonzero()
+        print("dif    : ", z[index])
+        print("inp1   : ", a[index])
+        print("inp2   : ", b[index])
+    return close
+
+parser = argparse.ArgumentParser()
+parser.add_argument('--local_rank', type=int, default=0)
+parser.add_argument('--apex', action='store_true')
+args = parser.parse_args()
+
+
+torch.manual_seed(2809)
+# Setup DDP
+torch.cuda.set_device(args.local_rank)
+device = torch.device('cuda:{}'.format(args.local_rank))
+
+torch.distributed.init_process_group(
+    'nccl',
+    init_method='env://',
+    rank=args.local_rank,
+)
+
+# Setup model
+if args.apex:
+    model = nn.Sequential(
+        nn.Conv2d(3, 6, 3, 1, 1),
+        ApexSyncBatchNorm(6)
+    )
+else:
+    model = nn.Sequential(
+        nn.Conv2d(3, 6, 3, 1, 1),
+        nn.SyncBatchNorm(6)
+    )
+
+# Setup reference model
+model_reference = nn.Sequential(
+    nn.Conv2d(3, 6, 3, 1, 1),
+    nn.BatchNorm2d(6)
+)
+
+with torch.no_grad():
+    model_reference[0].weight.copy_(model[0].weight)
+    model_reference[0].bias.copy_(model[0].bias)
+model_reference.to(device)
+
+model = model.to(device)
+model = DDP(model, device_ids=[args.local_rank], output_device=args.local_rank)
+
+global_batch_size = var_batch + 8
+# Create random data
+if args.local_rank == 0:
+    data = torch.randn(var_batch, 3, 8, 8, device=device, dtype=torch.float) * 50.0
+    grad = torch.randint(0, 10, (var_batch, 6, 8, 8), device=device, dtype=torch.float) / 10.0
+else:
+    data = torch.randn(8, 3, 8, 8, device=device)
+    grad = torch.randint(0, 10, (8, 6, 8, 8), device=device, dtype=torch.float) / 10.0
+
+data.requires_grad_()
+data.retain_grad = True
+
+weighted_gradient = True 
+
+# DDP forward/backward
+output = model(data)
+
+if weighted_gradient:
+    output.backward(grad * 2 / global_batch_size)
+else:
+    output.backward(grad / output.size(0))
+
+d_list = [torch.randn(8, 3, 8, 8, device=device) for i in range(int(os.environ['WORLD_SIZE']))]
+y_list = [torch.randn(8, 6, 8, 8, device=device) for i in range(int(os.environ['WORLD_SIZE']))]
+dgrad_list = [torch.randn(8, 3, 8, 8, device=device) for i in range(int(os.environ['WORLD_SIZE']))]
+grad_list = [torch.randn(8, 6, 8, 8, device=device) for i in range(int(os.environ['WORLD_SIZE']))]
+if args.local_rank == 0:
+    # placeholder, these random data will later be discarded.
+    torch.distributed.all_gather(d_list, torch.randn(8, 3, 8, 8, device=device))
+    torch.distributed.all_gather(y_list, torch.randn(8, 6, 8, 8, device=device))
+    torch.distributed.all_gather(dgrad_list, torch.randn(8, 3, 8, 8, device=device))
+    torch.distributed.all_gather(grad_list, torch.randn(8, 6, 8, 8, device=device))
+else:
+    torch.distributed.all_gather(d_list, data)
+    torch.distributed.all_gather(y_list, output)
+    torch.distributed.all_gather(dgrad_list, data.grad)
+    torch.distributed.all_gather(grad_list, grad)
+
+torch.distributed.barrier()
+
+if args.local_rank == 0:
+    ref_tensor = d_list[1:]
+    ref_tensor.insert(0, data)
+    assert(ref_tensor[0].equal(data))
+    ref_tensor = torch.cat(ref_tensor, 0)
+    ref_tensor = ref_tensor.detach()
+    ref_tensor.requires_grad_()
+    ref_tensor.retain_grad()
+
+    # Reference forward/backward
+    output_reference = model_reference(ref_tensor)
+    grad_tensor = grad_list[1:]
+    grad_tensor.insert(0, grad)
+    assert(grad_tensor[0].equal(grad))
+    grad_tensor = torch.cat(grad_tensor, 0)
+    if weighted_gradient:
+        output_reference.backward(grad_tensor / output_reference.size(0))
+    else:
+        output_reference.backward(grad_tensor / output_reference.size(0))
+
+    dgrad_tensor = dgrad_list[1:]
+    dgrad_tensor.insert(0, data.grad)
+    dgrad_tensor = torch.cat(dgrad_tensor, 0)
+    # check output
+    output_tensor = y_list[1:]
+    output_tensor.insert(0, output)
+    output_tensor = torch.cat(output_tensor, 0)
+    passed = True
+    passed = passed and compare("check output",
+          output_tensor,
+          output_reference)
+    # check stats
+    passed = passed and compare("check running mean failed",
+          model_reference[1].running_mean,
+          model.module[1].running_mean)
+    passed = passed and compare("check running var failed",
+          model_reference[1].running_var,
+          model.module[1].running_var)
+    passed = passed and compare("bn wgrad check failed!",
+          model_reference[1].weight.grad,
+          model.module[1].weight.grad, 1e-6)
+    passed = passed and compare("conv wgrad check failed!",
+          model_reference[0].weight.grad,
+          model.module[0].weight.grad)
+    # can't really compare dgrad directly, as we need to scale it to account for
+    # DDP
+    # passed = passed and compare("dgrad check failed!", ref_tensor.grad, dgrad_tensor)
+    if passed:
+      print("====SBN two gpu with different batches test passed")
+    else:
+      assert("*failed two gpu with different batches tests*")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/two_gpu_unit_test.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/two_gpu_unit_test.py
new file mode 100644
index 0000000000..ea89e7ae1a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/two_gpu_unit_test.py
@@ -0,0 +1,180 @@
+import torch
+import numpy as np
+import apex
+import syncbn
+import os
+import argparse
+import torch.optim as optim
+
+def compare(desc, inp1, inp2, error):
+    a = inp1.clone().detach().cpu().numpy()
+    b = inp2.clone().detach().cpu().numpy()
+    close = np.allclose(a,b, error, error)
+    if not close:
+        print(desc, close)
+        z = a - b
+        index = (np.abs(z) >= error + error * np.abs(b)).nonzero()
+        print("dif    : ", z[index])
+        print("inp1   : ", a[index])
+        print("inp2   : ", b[index])
+    return close
+
+feature_size = 10
+space_size = 40
+batch_size = 32
+
+
+from apex.parallel import DistributedDataParallel as DDP
+parser = argparse.ArgumentParser()
+parser.add_argument("--local_rank", default=0, type=int)
+parser.add_argument("--fp16", action='store_true', default=False)
+parser.add_argument("--fp64", action='store_true', default=False)
+args = parser.parse_args()
+args.world_size = int(os.environ['WORLD_SIZE'])
+torch.cuda.set_device(args.local_rank)
+torch.distributed.init_process_group(backend='nccl', init_method='env://')
+start = args.local_rank * batch_size//args.world_size
+finish = (args.local_rank + 1) * batch_size//args.world_size
+
+error = 1e-5
+dtype = np.float32
+if args.fp16:
+    error = 1e-3
+    dtype = np.float16
+elif args.fp64:
+    error = 1e-8
+    dtype = np.float64
+
+np.random.seed(18)
+inp = np.random.randn(batch_size, feature_size, space_size, space_size).astype(dtype)
+grad = np.random.randn(batch_size, feature_size, space_size, space_size).astype(dtype)
+weight = np.random.randn(feature_size).astype(dtype)
+bias = np.random.randn(feature_size).astype(dtype)
+
+
+type_tensor = torch.cuda.FloatTensor
+if args.fp16:
+    type_tensor = torch.cuda.HalfTensor
+if args.fp64:
+    type_tensor = torch.cuda.DoubleTensor
+
+ref_tensor = torch.cuda.DoubleTensor
+
+inp_t = type_tensor(inp)
+weight_t = type_tensor(weight)
+bias_t = type_tensor(bias)
+
+inp_r = ref_tensor(inp.transpose(1, 0, 2, 3).reshape(feature_size, -1))
+inp2_r = ref_tensor(inp)
+weight_r = ref_tensor(weight).view(-1, 1, 1)
+bias_r = ref_tensor(bias).view(-1, 1, 1)
+
+grad_output_t = type_tensor(grad)
+
+m = inp_r.mean(1)
+b_v = inp_r.var(1, unbiased=False)
+unb_v = inp_r.var(1, unbiased=True)
+
+eps = 1e-5
+
+mean, var_biased = syncbn.welford_mean_var(inp_t)
+inv_std = 1.0 / torch.sqrt(var_biased + eps)
+
+bn = torch.nn.BatchNorm2d(feature_size).cuda()
+bn.momentum = 1.0
+bn.weight.data = weight_t.clone()
+bn.bias.data = bias_t.clone()
+if args.fp16:
+    bn.half()
+if args.fp64:
+    bn.double()
+inp_bn = inp_t.clone().requires_grad_()
+grad_bn = grad_output_t.clone().detach()
+out_bn = bn(inp_bn)
+out_bn.backward(grad_bn)
+# compensating the averaging over processes done by DDP
+# in order to produce mathematically equivalent result
+# https://github.com/NVIDIA/apex/issues/134#issuecomment-458307368
+for param in bn.parameters():
+    param.grad = param.grad / args.world_size
+bn_opt = optim.SGD(bn.parameters(), lr=1.0)
+
+sbn = apex.parallel.SyncBatchNorm(feature_size).cuda()
+sbn.momentum = 1.0
+sbn.weight.data = weight_t.clone()
+sbn.bias.data = bias_t.clone()
+if args.fp16:
+    sbn.half()
+if args.fp64:
+    sbn.double()
+sbn = DDP(sbn)
+sbn_opt = optim.SGD(sbn.parameters(), lr=1.0)
+inp_sbn = inp_t.clone().requires_grad_()
+grad_sbn = grad_output_t.clone().detach()
+out_sbn = sbn(inp_sbn[start:finish])
+out_sbn.backward(grad_sbn[start:finish])
+
+count = [ space_size**2 * ( (i+1) * batch_size // args.world_size - i * batch_size // args.world_size ) for i in range(0, args.world_size)]
+count = torch.cuda.IntTensor(count)
+
+print("--- count : " , count)
+
+sbn_result = True
+bn_result = True
+
+if args.local_rank == 0:
+    sbn_result = compare("comparing mean: ", mean, m, error) and sbn_result
+    sbn_result = compare("comparing biased variance: ", var_biased, b_v, error) and sbn_result
+
+out = syncbn.batchnorm_forward(inp_t, mean, inv_std, weight_t, bias_t)
+out_r = weight_r * (inp2_r - m.view(-1, 1, 1)) * torch.rsqrt(b_v.view(-1,1,1) + eps) + bias_r
+
+if args.local_rank == 0:
+    sbn_result = compare("comparing output: ", out, out_r, error) and sbn_result
+    compare("comparing bn output: ", out_bn, out_r, error)
+
+grad_output_t = type_tensor(grad)
+
+grad_output_r = ref_tensor(grad.transpose(1, 0, 2, 3).reshape(feature_size, -1))
+grad_output2_r = ref_tensor(grad)
+
+grad_bias_r = grad_output_r.sum(1)
+grad_weight_r = ((inp2_r - m.view(-1, 1, 1)) * torch.rsqrt(b_v.view(-1,1,1) + eps) * grad_output2_r).transpose(1,0).contiguous().view(feature_size, -1).sum(1)
+
+sum_dy_r = grad_output_r.sum(1)
+mean_dy_r = grad_output_r.mean(1)
+mean_dy_xmu_r = ((inp2_r - m.view(-1, 1, 1)) * grad_output2_r).transpose(1,0).contiguous().view(feature_size, -1).mean(1)
+sum_dy_xmu_r = ((inp2_r - m.view(-1, 1, 1)) * grad_output2_r).transpose(1,0).contiguous().view(feature_size, -1).sum(1)
+
+grad_input_r = (grad_output2_r - mean_dy_r.view(-1, 1, 1) - (inp2_r - m.view(-1, 1, 1)) / (b_v.view(-1,1,1) + eps) * mean_dy_xmu_r.view(-1, 1, 1) ) * torch.rsqrt(b_v.view(-1,1,1) + eps) * weight_r.view(-1,1,1)
+
+sum_dy, sum_dy_xmu, grad_weight, grad_bias = syncbn.reduce_bn(grad_output_t, inp_t, mean, inv_std, weight_t)
+grad_input = syncbn.batchnorm_backward(grad_output_t, inp_t, mean, inv_std, weight_t, sum_dy, sum_dy_xmu, count)
+if args.local_rank == 0:
+    sbn_result = compare("comparing bias grad: ", grad_bias, grad_bias_r, error) and sbn_result
+    sbn_result = compare("comparing weight grad: ", grad_weight, grad_weight_r, error) and sbn_result
+    sbn_result = compare("comparing sum_dy grad: ", sum_dy, sum_dy_r, error) and sbn_result
+    sbn_result = compare("comparing sum_dy_xmu grad: ", sum_dy_xmu, sum_dy_xmu_r, error) and sbn_result
+    sbn_result = compare("comparing input grad: ", grad_input, grad_input_r, error) and sbn_result
+    compare("comparing bn input grad: ", inp_bn.grad, grad_input_r, error)
+
+if args.local_rank == 0:
+    sbn_result = compare("comparing running_mean: ", bn.running_mean.data, sbn.module.running_mean.data, error) and sbn_result
+    sbn_result = compare("comparing running_variance: ", bn.running_var.data, sbn.module.running_var.data, error) and sbn_result
+
+# execute by both
+compare("comparing layers output: ", out_bn[start:finish], out_sbn, error) and sbn_result
+compare("comparing layers grad_input: ", inp_bn.grad[start:finish], inp_sbn.grad[start:finish], error) and sbn_result
+
+bn_opt.step()
+sbn_opt.step()
+
+if args.local_rank == 0:
+    compare("comparing bn vs sbn bias: ", bn.bias, sbn.module.bias, error)
+    compare("comparing bn vs sbn weight: ", bn.weight, sbn.module.weight, error)
+
+
+if sbn_result:
+    print("====SBN two gpu passed tests")
+else:
+    print("*SBN two gpu failed*")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/unit_test.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/unit_test.sh
new file mode 100644
index 0000000000..2165f5ec06
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/distributed/synced_batchnorm/unit_test.sh
@@ -0,0 +1,8 @@
+python python_single_gpu_unit_test.py
+python single_gpu_unit_test.py
+python test_batchnorm1d.py
+python -m torch.distributed.launch --nproc_per_node=2 two_gpu_unit_test.py
+python -m torch.distributed.launch --nproc_per_node=2 two_gpu_unit_test.py --fp16
+python -m torch.distributed.launch --nproc_per_node=2 two_gpu_test_different_batch_size.py --apex
+#beware, you need a system with at least 4 gpus to test group_size<world_size
+#python -m torch.distributed.launch --nproc_per_node=4 test_groups.py --group_size=2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/docker_extension_builds/run.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/docker_extension_builds/run.sh
new file mode 100644
index 0000000000..3fb701058d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/apex/tests/docker_extension_builds/run.sh
@@ -0,0 +1,73 @@
+#!/bin/bash
+
+print_banner() {
+  printf "\n\n\n\e[30m\e[42m$1\e[0m\n\n\n\n"
+}
+
+print_green() {
+  printf "\e[30m\e[42m$1\e[0m\n"
+}
+
+print_red() {
+  printf "\e[30m\e[41m$1\e[0m\n"
+}
+
+images=(
+"gitlab-master.nvidia.com:5005/dl/dgx/pytorch:19.08-py3-devel"
+"gitlab-master.nvidia.com:5005/dl/dgx/pytorch:master-py3-devel"
+"pytorch/pytorch:nightly-devel-cuda10.0-cudnn7"
+"pytorch/pytorch:1.1.0-cuda10.0-cudnn7.5-devel"
+"pytorch/pytorch:1.0.1-cuda10.0-cudnn7-devel"
+"pytorch/pytorch:1.0-cuda10.0-cudnn7-devel"
+"pytorch/pytorch:nightly-devel-cuda9.2-cudnn7"
+)
+
+branch="master"
+
+# Associative array for exit codes
+declare -A exit_codes
+for image in images
+do
+  exit_codes[$image]="None"
+done
+
+for image in "${images[@]}"
+do
+  print_banner "$image"
+  set -x
+  docker pull $image
+  # Trying python setup.py install instead of pip install to ensure direct access to error codes.
+  # Maybe pip install would be ok too but this works.
+  docker run --runtime=nvidia --rm $image /bin/bash -c "yes | pip uninstall apex; yes | pip uninstall apex; git clone https://github.com/NVIDIA/apex.git; cd apex; git checkout $branch; set -e;  python setup.py install --cuda_ext --cpp_ext"
+  exit_code=$?
+  set +x
+  if [ $exit_code != 0 ]
+  then
+    print_red "Exit code: $exit_code"
+  else
+    print_green "Exit code: $exit_code"
+  fi
+  exit_codes[$image]=$exit_code
+done
+
+success=0
+for image in "${images[@]}"
+do
+  exit_code=${exit_codes[$image]}
+  if [ $exit_code != 0 ]
+  then
+    print_red "$image : $exit_code"
+    success=1
+  else
+    print_green "$image : $exit_code"
+  fi
+done
+
+if [ $success != 0 ]
+then
+  print_red "Overall status:  failure"
+else
+  print_green "Overall status:  success"
+fi
+
+exit $success
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/audio.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/audio.py
new file mode 100644
index 0000000000..916394f50d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/audio.py
@@ -0,0 +1,247 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import random
+import soundfile as sf
+
+import librosa
+import torch
+import numpy as np
+
+import sox
+
+
+def audio_from_file(file_path, offset=0, duration=0, trim=False, target_sr=16000):
+    audio = AudioSegment(file_path, target_sr=target_sr, int_values=False,
+                         offset=offset, duration=duration, trim=trim)
+
+    samples = torch.tensor(audio.samples, dtype=torch.float).cuda()
+    num_samples = torch.tensor(samples.shape[0]).int().cuda()
+    return (samples.unsqueeze(0), num_samples.unsqueeze(0))
+
+
+class AudioSegment(object):
+    """Monaural audio segment abstraction.
+
+    :param samples: Audio samples [num_samples x num_channels].
+    :type samples: ndarray.float32
+    :param sample_rate: Audio sample rate.
+    :type sample_rate: int
+    :raises TypeError: If the sample data type is not float or int.
+    """
+
+    def __init__(self, filename, target_sr=None, int_values=False, offset=0,
+                 duration=0, trim=False, trim_db=60):
+        """Create audio segment from samples.
+
+        Samples are convert float32 internally, with int scaled to [-1, 1].
+        Load a file supported by librosa and return as an AudioSegment.
+        :param filename: path of file to load
+        :param target_sr: the desired sample rate
+        :param int_values: if true, load samples as 32-bit integers
+        :param offset: offset in seconds when loading audio
+        :param duration: duration in seconds when loading audio
+        :return: numpy array of samples
+        """
+        with sf.SoundFile(filename, 'r') as f:
+            dtype = 'int32' if int_values else 'float32'
+            sample_rate = f.samplerate
+            if offset > 0:
+                f.seek(int(offset * sample_rate))
+            if duration > 0:
+                samples = f.read(int(duration * sample_rate), dtype=dtype)
+            else:
+                samples = f.read(dtype=dtype)
+        samples = samples.transpose()
+
+        samples = self._convert_samples_to_float32(samples)
+        if target_sr is not None and target_sr != sample_rate:
+            samples = librosa.core.resample(samples, sample_rate, target_sr)
+            sample_rate = target_sr
+        if trim:
+            samples, _ = librosa.effects.trim(samples, trim_db)
+        self._samples = samples
+        self._sample_rate = sample_rate
+        if self._samples.ndim >= 2:
+            self._samples = np.mean(self._samples, 1)
+
+    def __eq__(self, other):
+        """Return whether two objects are equal."""
+        if type(other) is not type(self):
+            return False
+        if self._sample_rate != other._sample_rate:
+            return False
+        if self._samples.shape != other._samples.shape:
+            return False
+        if np.any(self.samples != other._samples):
+            return False
+        return True
+
+    def __ne__(self, other):
+        """Return whether two objects are unequal."""
+        return not self.__eq__(other)
+
+    def __str__(self):
+        """Return human-readable representation of segment."""
+        return ("%s: num_samples=%d, sample_rate=%d, duration=%.2fsec, "
+                        "rms=%.2fdB" % (type(self), self.num_samples, self.sample_rate,
+                                                        self.duration, self.rms_db))
+
+    @staticmethod
+    def _convert_samples_to_float32(samples):
+        """Convert sample type to float32.
+
+        Audio sample type is usually integer or float-point.
+        Integers will be scaled to [-1, 1] in float32.
+        """
+        float32_samples = samples.astype('float32')
+        if samples.dtype in np.sctypes['int']:
+            bits = np.iinfo(samples.dtype).bits
+            float32_samples *= (1. / 2 ** (bits - 1))
+        elif samples.dtype in np.sctypes['float']:
+            pass
+        else:
+            raise TypeError("Unsupported sample type: %s." % samples.dtype)
+        return float32_samples
+
+    @property
+    def samples(self):
+        return self._samples.copy()
+
+    @property
+    def sample_rate(self):
+        return self._sample_rate
+
+    @property
+    def num_samples(self):
+        return self._samples.shape[0]
+
+    @property
+    def duration(self):
+        return self._samples.shape[0] / float(self._sample_rate)
+
+    @property
+    def rms_db(self):
+        mean_square = np.mean(self._samples ** 2)
+        return 10 * np.log10(mean_square)
+
+    def gain_db(self, gain):
+        self._samples *= 10. ** (gain / 20.)
+
+    def pad(self, pad_size, symmetric=False):
+        """Add zero padding to the sample.
+
+        The pad size is given in number of samples. If symmetric=True,
+        `pad_size` will be added to both sides. If false, `pad_size` zeros
+        will be added only to the end.
+        """
+        self._samples = np.pad(self._samples,
+                               (pad_size if symmetric else 0, pad_size),
+                               mode='constant')
+
+    def subsegment(self, start_time=None, end_time=None):
+        """Cut the AudioSegment between given boundaries.
+
+        Note that this is an in-place transformation.
+        :param start_time: Beginning of subsegment in seconds.
+        :type start_time: float
+        :param end_time: End of subsegment in seconds.
+        :type end_time: float
+        :raise ValueError: If start_time or end_time is incorrectly set, e.g. out
+                                             of bounds in time.
+        """
+        start_time = 0.0 if start_time is None else start_time
+        end_time = self.duration if end_time is None else end_time
+        if start_time < 0.0:
+            start_time = self.duration + start_time
+        if end_time < 0.0:
+            end_time = self.duration + end_time
+        if start_time < 0.0:
+            raise ValueError("The slice start position (%f s) is out of "
+                             "bounds." % start_time)
+        if end_time < 0.0:
+            raise ValueError("The slice end position (%f s) is out of bounds." %
+                             end_time)
+        if start_time > end_time:
+            raise ValueError("The slice start position (%f s) is later than "
+                             "the end position (%f s)." % (start_time, end_time))
+        if end_time > self.duration:
+            raise ValueError("The slice end position (%f s) is out of bounds "
+                             "(> %f s)" % (end_time, self.duration))
+        start_sample = int(round(start_time * self._sample_rate))
+        end_sample = int(round(end_time * self._sample_rate))
+        self._samples = self._samples[start_sample:end_sample]
+
+
+class Perturbation:
+    def __init__(self, p=0.1, rng=None):
+        self.p = p
+        self._rng = random.Random() if rng is None else rng
+
+    def maybe_apply(self, segment, sample_rate=None):
+        if self._rng.random() < self.p:
+            self(segment, sample_rate)
+
+
+class SpeedPerturbation(Perturbation):
+    def __init__(self, min_rate=0.85, max_rate=1.15, discrete=False, p=0.1, rng=None):
+        super(SpeedPerturbation, self).__init__(p, rng)
+        assert 0 < min_rate < max_rate
+        self.min_rate = min_rate
+        self.max_rate = max_rate
+        self.discrete = discrete
+
+    def __call__(self, data, sample_rate):
+        if self.discrete:
+            rate = np.random.choice([self.min_rate, None, self.max_rate])
+        else:
+            rate = self._rng.uniform(self.min_rate, self.max_rate)
+
+        if rate is not None:
+            data._samples = sox.Transformer().speed(factor=rate).build_array(
+                input_array=data._samples, sample_rate_in=sample_rate)
+
+
+class GainPerturbation(Perturbation):
+    def __init__(self, min_gain_dbfs=-10, max_gain_dbfs=10, p=0.1, rng=None):
+        super(GainPerturbation, self).__init__(p, rng)
+        self._rng = random.Random() if rng is None else rng
+        self._min_gain_dbfs = min_gain_dbfs
+        self._max_gain_dbfs = max_gain_dbfs
+
+    def __call__(self, data, sample_rate=None):
+        del sample_rate
+        gain = self._rng.uniform(self._min_gain_dbfs, self._max_gain_dbfs)
+        data._samples = data._samples * (10. ** (gain / 20.))
+
+
+class ShiftPerturbation(Perturbation):
+    def __init__(self, min_shift_ms=-5.0, max_shift_ms=5.0, p=0.1, rng=None):
+        super(ShiftPerturbation, self).__init__(p, rng)
+        self._min_shift_ms = min_shift_ms
+        self._max_shift_ms = max_shift_ms
+
+    def __call__(self, data, sample_rate):
+        shift_ms = self._rng.uniform(self._min_shift_ms, self._max_shift_ms)
+        if abs(shift_ms) / 1000 > data.duration:
+            # TODO: do something smarter than just ignore this condition
+            return
+        shift_samples = int(shift_ms * data.sample_rate // 1000)
+        # print("DEBUG: shift:", shift_samples)
+        if shift_samples < 0:
+            data._samples[-shift_samples:] = data._samples[:shift_samples]
+            data._samples[:-shift_samples] = 0
+        elif shift_samples > 0:
+            data._samples[:-shift_samples] = data._samples[shift_samples:]
+            data._samples[-shift_samples:] = 0
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/__init__.py
new file mode 100644
index 0000000000..ff800034d6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/__init__.py
@@ -0,0 +1,13 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/data_loader.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/data_loader.py
new file mode 100644
index 0000000000..99c06c695a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/data_loader.py
@@ -0,0 +1,158 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import os
+import math
+import numpy as np
+import torch.distributed as dist
+from .iterator import DaliJasperIterator, SyntheticDataIterator
+from .pipeline import DaliPipeline
+from common.helpers import print_once
+
+
+def _parse_json(json_path: str, start_label=0, predicate=lambda json: True):
+    """
+    Parses json file to the format required by DALI
+    Args:
+        json_path: path to json file
+        start_label: the label, starting from which DALI will assign consecutive int numbers to every transcript
+        predicate: function, that accepts a sample descriptor (i.e. json dictionary) as an argument.
+                   If the predicate for a given sample returns True, it will be included in the dataset.
+
+    Returns:
+        output_files: dictionary, that maps file name to label assigned by DALI
+        transcripts: dictionary, that maps label assigned by DALI to the transcript
+    """
+    import json
+    global cnt
+    with open(json_path) as f:
+        librispeech_json = json.load(f)
+    output_files = {}
+    transcripts = {}
+    curr_label = start_label
+    for original_sample in librispeech_json:
+        if not predicate(original_sample):
+            continue
+        transcripts[curr_label] = original_sample['transcript']
+        output_files[original_sample['files'][-1]['fname']] = curr_label
+        curr_label += 1
+    return output_files, transcripts
+
+
+def _dict_to_file(dict: dict, filename: str):
+    with open(filename, "w") as f:
+        for key, value in dict.items():
+            f.write("{} {}\n".format(key, value))
+
+
+class DaliDataLoader:
+    """
+    DataLoader is the main entry point to the data preprocessing pipeline.
+    To use, create an object and then just iterate over `data_iterator`.
+    DataLoader will do the rest for you.
+    Example:
+        data_layer = DataLoader(DaliTrainPipeline, path, json, bs, ngpu)
+        data_it = data_layer.data_iterator
+        for data in data_it:
+            print(data)  # Here's your preprocessed data
+
+    Args:
+        device_type: Which device to use for preprocessing. Choose: "cpu", "gpu"
+        pipeline_type: Choose: "train", "val", "synth"
+    """
+
+    def __init__(self, gpu_id, dataset_path: str, config_data: dict, config_features: dict, json_names: list,
+                 symbols: list, batch_size: int, pipeline_type: str, grad_accumulation_steps: int = 1,
+                 synth_iters_per_epoch: int = 544, device_type: str = "gpu"):
+        import torch
+        self.batch_size = batch_size
+        self.grad_accumulation_steps = grad_accumulation_steps
+        self.drop_last = (pipeline_type == 'train')
+        self.device_type = device_type
+        pipeline_type = self._parse_pipeline_type(pipeline_type)
+        if pipeline_type == "synth":
+            self._dali_data_iterator = self._init_synth_iterator(self.batch_size, config_features['nfilt'],
+                                                                 iters_per_epoch=synth_iters_per_epoch,
+                                                                 ngpus=torch.distributed.get_world_size())
+        else:
+            self._dali_data_iterator = self._init_iterator(gpu_id=gpu_id, dataset_path=dataset_path,
+                                                           config_data=config_data,
+                                                           config_features=config_features,
+                                                           json_names=json_names, symbols=symbols,
+                                                           train_pipeline=pipeline_type == "train")
+
+    def _init_iterator(self, gpu_id, dataset_path, config_data, config_features, json_names: list, symbols: list,
+                       train_pipeline: bool):
+        """
+        Returns data iterator. Data underneath this operator is preprocessed within Dali
+        """
+
+        def hash_list_of_strings(li):
+            return str(abs(hash(''.join(li))))
+
+        output_files, transcripts = {}, {}
+        max_duration = config_data['max_duration']
+        for jname in json_names:
+            of, tr = _parse_json(jname if jname[0] == '/' else os.path.join(dataset_path, jname), len(output_files),
+                                 predicate=lambda json: json['original_duration'] <= max_duration)
+            output_files.update(of)
+            transcripts.update(tr)
+        file_list_path = os.path.join("/tmp", "jasper_dali.file_list." + hash_list_of_strings(json_names))
+        _dict_to_file(output_files, file_list_path)
+        self.dataset_size = len(output_files)
+        print_once(f"Dataset read by DALI. Number of samples: {self.dataset_size}")
+
+        pipeline = DaliPipeline.from_config(config_data=config_data, config_features=config_features, device_id=gpu_id,
+                                            file_root=dataset_path, file_list=file_list_path,
+                                            device_type=self.device_type, batch_size=self.batch_size,
+                                            train_pipeline=train_pipeline)
+
+        return DaliJasperIterator([pipeline], transcripts=transcripts, symbols=symbols, batch_size=self.batch_size,
+                                  reader_name="file_reader", train_iterator=train_pipeline)
+
+    def _init_synth_iterator(self, batch_size, nfeatures, iters_per_epoch, ngpus):
+        self.dataset_size = ngpus * iters_per_epoch * batch_size
+        return SyntheticDataIterator(batch_size, nfeatures, regenerate=True)
+
+    @staticmethod
+    def _parse_pipeline_type(pipeline_type):
+        pipe = pipeline_type.lower()
+        assert pipe in ("train", "val", "synth"), 'Invalid pipeline type (choices: "train", "val", "synth").'
+        return pipe
+
+    def _shard_size(self):
+        """
+        Total number of samples handled by a single GPU in a single epoch.
+        """
+        world_size = dist.get_world_size() if dist.is_initialized() else 1
+        if self.drop_last:
+            divisor = world_size * self.batch_size * self.grad_accumulation_steps
+            return self.dataset_size // divisor * divisor // world_size
+        else:
+            return int(math.ceil(self.dataset_size / world_size))
+
+    def __len__(self):
+        """
+        Number of batches handled by each GPU.
+        """
+        if self.drop_last:
+            assert self._shard_size() % self.batch_size == 0, f'{self._shard_size()} {self.batch_size}'
+
+        return int(math.ceil(self._shard_size() / self.batch_size))
+
+    def data_iterator(self):
+        return self._dali_data_iterator
+
+    def __iter__(self):
+        return self._dali_data_iterator
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/iterator.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/iterator.py
new file mode 100644
index 0000000000..ea11013500
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/iterator.py
@@ -0,0 +1,161 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import torch
+import torch.distributed as dist
+import numpy as np
+from common.helpers import print_once
+from common.text import _clean_text, punctuation_map
+
+
+def normalize_string(s, symbols, punct_map):
+    """
+    Normalizes string.
+    Example:
+        'call me at 8:00 pm!' -> 'call me at eight zero pm'
+    """
+    labels = set(symbols)
+    try:
+        text = _clean_text(s, ["english_cleaners"], punct_map).strip()
+        return ''.join([tok for tok in text if all(t in labels for t in tok)])
+    except Exception as e:
+        print_once("WARNING: Normalizing failed: {s} {e}")
+
+
+class DaliJasperIterator(object):
+    """
+    Returns batches of data for Jasper training:
+    preprocessed_signal, preprocessed_signal_length, transcript, transcript_length
+
+    This iterator is not meant to be the entry point to Dali processing pipeline.
+    Use DataLoader instead.
+    """
+
+    def __init__(self, dali_pipelines, transcripts, symbols, batch_size, reader_name, train_iterator: bool):
+        self.transcripts = transcripts
+        self.symbols = symbols
+        self.batch_size = batch_size
+        from nvidia.dali.plugin.pytorch import DALIGenericIterator
+        from nvidia.dali.plugin.base_iterator import LastBatchPolicy
+
+        # in train pipeline shard_size is set to divisable by batch_size, so PARTIAL policy is safe
+        self.dali_it = DALIGenericIterator(
+            dali_pipelines, ["audio", "label", "audio_shape"], reader_name=reader_name,
+            dynamic_shape=True, auto_reset=True, last_batch_policy=LastBatchPolicy.PARTIAL)
+
+    @staticmethod
+    def _str2list(s: str):
+        """
+        Returns list of floats, that represents given string.
+        '0.' denotes separator
+        '1.' denotes 'a'
+        '27.' denotes "'"
+        Assumes, that the string is lower case.
+        """
+        list = []
+        for c in s:
+            if c == "'":
+                list.append(27.)
+            else:
+                list.append(max(0., ord(c) - 96.))
+        return list
+
+    @staticmethod
+    def _pad_lists(lists: list, pad_val=0):
+        """
+        Pads lists, so that all have the same size.
+        Returns list with actual sizes of corresponding input lists
+        """
+        max_length = 0
+        sizes = []
+        for li in lists:
+            sizes.append(len(li))
+            max_length = max_length if len(li) < max_length else len(li)
+        for li in lists:
+            li += [pad_val] * (max_length - len(li))
+        return sizes
+
+    def _gen_transcripts(self, labels, normalize_transcripts: bool = True):
+        """
+        Generate transcripts in format expected by NN
+        """
+        lists = [
+            self._str2list(normalize_string(self.transcripts[lab.item()], self.symbols, punctuation_map(self.symbols)))
+            for lab in labels
+        ] if normalize_transcripts else [self._str2list(self.transcripts[lab.item()]) for lab in labels]
+        sizes = self._pad_lists(lists)
+        return torch.tensor(lists).cuda(), torch.tensor(sizes, dtype=torch.int32).cuda()
+
+    def __next__(self):
+        data = self.dali_it.__next__()
+        transcripts, transcripts_lengths = self._gen_transcripts(data[0]["label"])
+        return data[0]["audio"], data[0]["audio_shape"][:, 1], transcripts, transcripts_lengths
+
+    def next(self):
+        return self.__next__()
+
+    def __iter__(self):
+        return self
+
+
+# TODO: refactor
+class SyntheticDataIterator(object):
+    def __init__(self, batch_size, nfeatures, feat_min=-5., feat_max=0., txt_min=0., txt_max=23., feat_lens_max=1760,
+                 txt_lens_max=231, regenerate=False):
+        """
+        Args:
+            batch_size
+            nfeatures: number of features for melfbanks
+            feat_min: minimum value in `feat` tensor, used for randomization
+            feat_max: maximum value in `feat` tensor, used for randomization
+            txt_min: minimum value in `txt` tensor, used for randomization
+            txt_max: maximum value in `txt` tensor, used for randomization
+            regenerate: If True, regenerate random tensors for every iterator step.
+                        If False, generate them only at start.
+        """
+        self.batch_size = batch_size
+        self.nfeatures = nfeatures
+        self.feat_min = feat_min
+        self.feat_max = feat_max
+        self.feat_lens_max = feat_lens_max
+        self.txt_min = txt_min
+        self.txt_max = txt_max
+        self.txt_lens_max = txt_lens_max
+        self.regenerate = regenerate
+
+        if not self.regenerate:
+            self.feat, self.feat_lens, self.txt, self.txt_lens = self._generate_sample()
+
+    def _generate_sample(self):
+        feat = (self.feat_max - self.feat_min) * np.random.random_sample(
+            (self.batch_size, self.nfeatures, self.feat_lens_max)) + self.feat_min
+        feat_lens = np.random.randint(0, int(self.feat_lens_max) - 1, size=self.batch_size)
+        txt = (self.txt_max - self.txt_min) * np.random.random_sample(
+            (self.batch_size, self.txt_lens_max)) + self.txt_min
+        txt_lens = np.random.randint(0, int(self.txt_lens_max) - 1, size=self.batch_size)
+        return torch.Tensor(feat).cuda(), \
+               torch.Tensor(feat_lens).cuda(), \
+               torch.Tensor(txt).cuda(), \
+               torch.Tensor(txt_lens).cuda()
+
+    def __next__(self):
+        if self.regenerate:
+            return self._generate_sample()
+        return self.feat, self.feat_lens, self.txt, self.txt_lens
+
+    def next(self):
+        return self.__next__()
+
+    def __iter__(self):
+        return self
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/pipeline.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/pipeline.py
new file mode 100644
index 0000000000..bdbb26b40c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dali/pipeline.py
@@ -0,0 +1,366 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import nvidia.dali as dali
+import nvidia.dali.fn as fn
+import nvidia.dali.types as types
+import multiprocessing
+import numpy as np
+import torch
+import math
+import itertools
+
+
+class DaliPipeline():
+    def __init__(self, *,
+                 train_pipeline: bool,  # True if train pipeline, False if validation pipeline
+                 device_id,
+                 num_threads,
+                 batch_size,
+                 file_root: str,
+                 file_list: str,
+                 sample_rate,
+                 discrete_resample_range: bool,
+                 resample_range: list,
+                 window_size,
+                 window_stride,
+                 nfeatures,
+                 nfft,
+                 frame_splicing_factor,
+                 dither_coeff,
+                 silence_threshold,
+                 preemph_coeff,
+                 pad_align,
+                 max_duration,
+                 mask_time_num_regions,
+                 mask_time_min,
+                 mask_time_max,
+                 mask_freq_num_regions,
+                 mask_freq_min,
+                 mask_freq_max,
+                 mask_both_num_regions,
+                 mask_both_min_time,
+                 mask_both_max_time,
+                 mask_both_min_freq,
+                 mask_both_max_freq,
+                 preprocessing_device="gpu",
+                 is_triton_pipeline=False):
+        self._dali_init_log(locals())
+
+        if torch.distributed.is_initialized():
+            shard_id = torch.distributed.get_rank()
+            n_shards = torch.distributed.get_world_size()
+        else:
+            shard_id = 0
+            n_shards = 1
+
+        self.preprocessing_device = preprocessing_device.lower()
+        assert self.preprocessing_device == "cpu" or self.preprocessing_device == "gpu", \
+            "Incorrect preprocessing device. Please choose either 'cpu' or 'gpu'"
+        self.frame_splicing_factor = frame_splicing_factor
+
+        # TODO(janton): Implement this
+        assert frame_splicing_factor == 1, "Frame splicing is not yet implemented"
+
+        self.resample_range = resample_range
+        self.discrete_resample_range = discrete_resample_range
+
+        self.train = train_pipeline
+        self.sample_rate = sample_rate
+        self.dither_coeff = dither_coeff
+        self.nfeatures = nfeatures
+        self.max_duration = max_duration
+        self.mask_params = {
+            'time_num_regions': mask_time_num_regions,
+            'time_min': mask_time_min,
+            'time_max': mask_time_max,
+            'freq_num_regions': mask_freq_num_regions,
+            'freq_min': mask_freq_min,
+            'freq_max': mask_freq_max,
+            'both_num_regions': mask_both_num_regions,
+            'both_min_time': mask_both_min_time,
+            'both_max_time': mask_both_max_time,
+            'both_min_freq': mask_both_min_freq,
+            'both_max_freq': mask_both_max_freq,
+        }
+        self.do_remove_silence = True if silence_threshold is not None else False
+
+        @dali.pipeline_def
+        def dali_jasper_pipe():
+            if is_triton_pipeline:
+                assert not self.train, "Pipeline for Triton shall be a validation pipeline"
+                if torch.distributed.is_initialized():
+                    raise RuntimeError(
+                        "You're creating Triton pipeline, using multi-process mode. Please use single-process mode.")
+                encoded, label = fn.external_source(device="cpu", name="DALI_INPUT_0", no_copy=True)
+            else:
+                encoded, label = fn.readers.file(device="cpu", name="file_reader",
+                                                 file_root=file_root, file_list=file_list, shard_id=shard_id,
+                                                 num_shards=n_shards, shuffle_after_epoch=train_pipeline)
+
+            speed_perturbation_coeffs = None
+            if resample_range is not None:
+                if discrete_resample_range:
+                    values = [self.resample_range[0], 1.0, self.resample_range[1]]
+                    speed_perturbation_coeffs = fn.random.uniform(device="cpu", values=values)
+                else:
+                    speed_perturbation_coeffs = fn.random.uniform(device="cpu", range=resample_range)
+
+            if self.train and speed_perturbation_coeffs is not None:
+                dec_sample_rate_arg = speed_perturbation_coeffs * self.sample_rate
+            elif resample_range is None:
+                dec_sample_rate_arg = self.sample_rate
+            else:
+                dec_sample_rate_arg = None
+
+            audio, _ = fn.decoders.audio(encoded, sample_rate=dec_sample_rate_arg, dtype=types.FLOAT, downmix=True)
+
+            if self.do_remove_silence:
+                begin, length = fn.nonsilent_region(audio, cutoff_db=silence_threshold)
+                audio = fn.slice(audio, begin, length, axes=[0])
+
+            # Max duration drop is performed at DataLayer stage
+
+            if self.preprocessing_device == "gpu":
+                audio = audio.gpu()
+
+            if self.dither_coeff != 0.:
+                audio = audio + fn.random.normal(audio) * self.dither_coeff
+
+            audio = fn.preemphasis_filter(audio, preemph_coeff=preemph_coeff)
+
+            spec = fn.spectrogram(audio, nfft=nfft,
+                                  window_length=window_size * sample_rate, window_step=window_stride * sample_rate)
+
+            mel_spec = fn.mel_filter_bank(spec, sample_rate=sample_rate, nfilter=self.nfeatures, normalize=True)
+
+            log_features = fn.to_decibels(mel_spec, multiplier=np.log(10), reference=1.0, cutoff_db=math.log(1e-20))
+
+            log_features_len = fn.shapes(log_features)
+            if self.frame_splicing_factor != 1:
+                log_features_len = self._div_ceil(log_features_len, self.frame_splicing_factor)
+
+            log_features = fn.normalize(log_features, axes=[1])
+            log_features = fn.pad(log_features, axes=[1], fill_value=0, align=pad_align)
+
+            if self.train and self._do_spectrogram_masking():
+                anchors, shapes  = fn.external_source(source=self._cutouts_generator, num_outputs=2, cycle=True)
+                log_features = fn.erase(log_features, anchor=anchors, shape=shapes, axes=[0, 1], fill_value=0,
+                                        normalized_anchor=True)
+
+            # When modifying DALI pipeline returns, make sure you update `output_map` in DALIGenericIterator invocation
+            return log_features.gpu(), label.gpu(), log_features_len.gpu()
+
+        self.pipe_handle = dali_jasper_pipe(batch_size=batch_size, num_threads=num_threads, device_id=device_id)
+
+    def get_pipeline(self):
+        return self.pipe_handle
+
+    @classmethod
+    def from_config(cls, train_pipeline: bool, device_id, batch_size, file_root: str, file_list: str, config_data: dict,
+                    config_features: dict, device_type: str = "gpu", do_resampling: bool = True,
+                    num_cpu_threads=multiprocessing.cpu_count()):
+
+        max_duration = config_data['max_duration']
+        sample_rate = config_data['sample_rate']
+        silence_threshold = -60 if config_data['trim_silence'] else None
+
+        # TODO Take into account resampling probablity
+        # TODO     config_features['speed_perturbation']['p']
+
+        if do_resampling and config_data['speed_perturbation'] is not None:
+            resample_range = [config_data['speed_perturbation']['min_rate'],
+                              config_data['speed_perturbation']['max_rate']]
+            discrete_resample_range = config_data['speed_perturbation']['discrete']
+        else:
+            resample_range = None
+            discrete_resample_range = False
+
+        window_size = config_features['window_size']
+        window_stride = config_features['window_stride']
+        nfeatures = config_features['n_filt']
+        nfft = config_features['n_fft']
+        frame_splicing_factor = config_features['frame_splicing']
+        dither_coeff = config_features['dither']
+        pad_align = config_features['pad_align']
+        pad_to_max_duration = config_features['pad_to_max_duration']
+        assert not pad_to_max_duration, "Padding to max duration currently not supported in DALI"
+        preemph_coeff = .97
+
+        config_spec = config_features['spec_augment']
+        if config_spec is not None:
+            mask_time_num_regions = config_spec['time_masks']
+            mask_time_min = config_spec['min_time']
+            mask_time_max = config_spec['max_time']
+            mask_freq_num_regions = config_spec['freq_masks']
+            mask_freq_min = config_spec['min_freq']
+            mask_freq_max = config_spec['max_freq']
+        else:
+            mask_time_num_regions = 0
+            mask_time_min = 0
+            mask_time_max = 0
+            mask_freq_num_regions = 0
+            mask_freq_min = 0
+            mask_freq_max = 0
+
+        config_cutout = config_features['cutout_augment']
+        if config_cutout is not None:
+            mask_both_num_regions = config_cutout['masks']
+            mask_both_min_time = config_cutout['min_time']
+            mask_both_max_time = config_cutout['max_time']
+            mask_both_min_freq = config_cutout['min_freq']
+            mask_both_max_freq = config_cutout['max_freq']
+        else:
+            mask_both_num_regions = 0
+            mask_both_min_time = 0
+            mask_both_max_time = 0
+            mask_both_min_freq = 0
+            mask_both_max_freq = 0
+
+        inst = cls(train_pipeline=train_pipeline,
+                   device_id=device_id,
+                   preprocessing_device=device_type,
+                   num_threads=num_cpu_threads,
+                   batch_size=batch_size,
+                   file_root=file_root,
+                   file_list=file_list,
+                   sample_rate=sample_rate,
+                   discrete_resample_range=discrete_resample_range,
+                   resample_range=resample_range,
+                   window_size=window_size,
+                   window_stride=window_stride,
+                   nfeatures=nfeatures,
+                   nfft=nfft,
+                   frame_splicing_factor=frame_splicing_factor,
+                   dither_coeff=dither_coeff,
+                   silence_threshold=silence_threshold,
+                   preemph_coeff=preemph_coeff,
+                   pad_align=pad_align,
+                   max_duration=max_duration,
+                   mask_time_num_regions=mask_time_num_regions,
+                   mask_time_min=mask_time_min,
+                   mask_time_max=mask_time_max,
+                   mask_freq_num_regions=mask_freq_num_regions,
+                   mask_freq_min=mask_freq_min,
+                   mask_freq_max=mask_freq_max,
+                   mask_both_num_regions=mask_both_num_regions,
+                   mask_both_min_time=mask_both_min_time,
+                   mask_both_max_time=mask_both_max_time,
+                   mask_both_min_freq=mask_both_min_freq,
+                   mask_both_max_freq=mask_both_max_freq)
+        return inst.get_pipeline()
+
+    @staticmethod
+    def _dali_init_log(args: dict):
+        if (not torch.distributed.is_initialized() or (
+                torch.distributed.is_initialized() and torch.distributed.get_rank() == 0)):  # print once
+            max_len = max([len(ii) for ii in args.keys()])
+            fmt_string = '\t%' + str(max_len) + 's : %s'
+            print('Initializing DALI with parameters:')
+            for keyPair in sorted(args.items()):
+                print(fmt_string % keyPair)
+
+    @staticmethod
+    def _div_ceil(dividend, divisor):
+        return (dividend + (divisor - 1)) // divisor
+
+    def _do_spectrogram_masking(self):
+        return self.mask_params['time_num_regions'] > 0 or self.mask_params['freq_num_regions'] > 0 or \
+               self.mask_params['both_num_regions'] > 0
+
+    @staticmethod
+    def _interleave_lists(*lists):
+        """
+        [*, **, ***], [1, 2, 3], [a, b, c] -> [*, 1, a, **, 2, b, ***, 3, c]
+        Returns:
+             iterator over interleaved list
+        """
+        assert all((len(lists[0]) == len(test_l) for test_l in lists)), "All lists have to have the same length"
+        return itertools.chain(*zip(*lists))
+
+    def _generate_cutouts(self):
+        """
+        Returns:
+            Generates anchors and shapes of the cutout regions.
+            Single call generates one batch of data.
+            The output shall be passed to DALI's Erase operator
+            anchors = [f0 t0 f1 t1 ...]
+            shapes = [f0w t0h f1w t1h ...]
+        """
+        MAX_TIME_DIMENSION = 20 * 16000
+        freq_anchors = np.random.random(self.mask_params['freq_num_regions'])
+        time_anchors = np.random.random(self.mask_params['time_num_regions'])
+        both_anchors_freq = np.random.random(self.mask_params['both_num_regions'])
+        both_anchors_time = np.random.random(self.mask_params['both_num_regions'])
+        anchors = []
+        for anch in freq_anchors:
+            anchors.extend([anch, 0])
+        for anch in time_anchors:
+            anchors.extend([0, anch])
+        for t, f in zip(both_anchors_time, both_anchors_freq):
+            anchors.extend([f, t])
+
+        shapes = []
+        shapes.extend(
+            self._interleave_lists(
+                np.random.randint(self.mask_params['freq_min'], self.mask_params['freq_max'] + 1,
+                                  self.mask_params['freq_num_regions']),
+                # XXX: Here, a time dimension of the spectrogram shall be passed.
+                #      However, in DALI ArgumentInput can't come from GPU.
+                #      So we leave the job for Erase (masking operator) to get it together.
+                [int(MAX_TIME_DIMENSION)] * self.mask_params['freq_num_regions']
+            )
+        )
+        shapes.extend(
+            self._interleave_lists(
+                [self.nfeatures] * self.mask_params['time_num_regions'],
+                np.random.randint(self.mask_params['time_min'], self.mask_params['time_max'] + 1,
+                                  self.mask_params['time_num_regions'])
+            )
+        )
+        shapes.extend(
+            self._interleave_lists(
+                np.random.randint(self.mask_params['both_min_freq'], self.mask_params['both_max_freq'] + 1,
+                                  self.mask_params['both_num_regions']),
+                np.random.randint(self.mask_params['both_min_time'], self.mask_params['both_max_time'] + 1,
+                                  self.mask_params['both_num_regions'])
+            )
+        )
+        return anchors, shapes
+
+    def _cutouts_generator(self):
+        """
+        Generator, that wraps cutouts creation in order to randomize inputs
+        and allow passing them to DALI's ExternalSource operator
+        """
+
+        def tuples2list(tuples: list):
+            """
+            [(a, b), (c, d)] -> [[a, c], [b, d]]
+            """
+            return map(list, zip(*tuples))
+
+        [anchors, shapes] = tuples2list([self._generate_cutouts() for _ in range(self.pipe_handle.max_batch_size)])
+        yield np.array(anchors, dtype=np.float32), np.array(shapes, dtype=np.float32)
+
+class DaliTritonPipeline(DaliPipeline):
+    def __init__(self, **kwargs):
+        kwargs['is_triton_pipeline'] = True
+        super().__init__(**kwargs)
+
+def serialize_dali_triton_pipeline(output_path: str, config_data: dict, config_features: dict):
+    pipe = DaliTritonPipeline.from_config(train_pipeline=False, device_id=-1, batch_size=-1, file_root=None,
+                                          file_list=None, config_data=config_data, config_features=config_features,
+                                          do_resampling=False, num_cpu_threads=-1)
+    pipe.serialize(filename=output_path)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dataset.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dataset.py
new file mode 100644
index 0000000000..daf4070d4d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/dataset.py
@@ -0,0 +1,237 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import json
+from pathlib import Path
+
+import numpy as np
+
+import torch
+from torch.utils.data import Dataset, DataLoader
+from torch.utils.data.distributed import DistributedSampler
+
+from .audio import (audio_from_file, AudioSegment, GainPerturbation,
+                    ShiftPerturbation, SpeedPerturbation)
+from .text import _clean_text, punctuation_map
+
+
+def normalize_string(s, labels, punct_map):
+    """Normalizes string.
+
+    Example:
+        'call me at 8:00 pm!' -> 'call me at eight zero pm'
+    """
+    labels = set(labels)
+    try:
+        text = _clean_text(s, ["english_cleaners"], punct_map).strip()
+        return ''.join([tok for tok in text if all(t in labels for t in tok)])
+    except:
+        print(f"WARNING: Normalizing failed: {s}")
+        return None
+
+
+class FilelistDataset(Dataset):
+    def __init__(self, filelist_fpath):
+        self.samples = [line.strip() for line in open(filelist_fpath, 'r')]
+
+    def __len__(self):
+        return len(self.samples)
+
+    def __getitem__(self, index):
+        audio, audio_len = audio_from_file(self.samples[index])
+        return (audio.squeeze(0), audio_len, torch.LongTensor([0]),
+                torch.LongTensor([0]))
+
+
+class SingleAudioDataset(FilelistDataset):
+    def __init__(self, audio_fpath):
+        self.samples = [audio_fpath]
+
+
+class AudioDataset(Dataset):
+    def __init__(self, data_dir, manifest_fpaths, labels,
+                 sample_rate=16000, min_duration=0.1, max_duration=float("inf"),
+                 pad_to_max_duration=False, max_utts=0, normalize_transcripts=True,
+                 sort_by_duration=False, trim_silence=False,
+                 speed_perturbation=None, gain_perturbation=None,
+                 shift_perturbation=None, ignore_offline_speed_perturbation=False):
+        """Loads audio, transcript and durations listed in a .json file.
+
+        Args:
+            data_dir: absolute path to dataset folder
+            manifest_filepath: relative path from dataset folder
+                to manifest json as described above. Can be coma-separated paths.
+            labels (str): all possible output symbols
+            min_duration (int): skip audio shorter than threshold
+            max_duration (int): skip audio longer than threshold
+            pad_to_max_duration (bool): pad all sequences to max_duration
+            max_utts (int): limit number of utterances
+            normalize_transcripts (bool): normalize transcript text
+            sort_by_duration (bool): sort sequences by increasing duration
+            trim_silence (bool): trim leading and trailing silence from audio
+            ignore_offline_speed_perturbation (bool): use precomputed speed perturbation
+
+        Returns:
+            tuple of Tensors
+        """
+        self.data_dir = data_dir
+        self.labels = labels
+        self.labels_map = dict([(labels[i], i) for i in range(len(labels))])
+        self.punctuation_map = punctuation_map(labels)
+        self.blank_index = len(labels)
+
+        self.pad_to_max_duration = pad_to_max_duration
+
+        self.sort_by_duration = sort_by_duration
+        self.max_utts = max_utts
+        self.normalize_transcripts = normalize_transcripts
+        self.ignore_offline_speed_perturbation = ignore_offline_speed_perturbation
+
+        self.min_duration = min_duration
+        self.max_duration = max_duration
+        self.trim_silence = trim_silence
+        self.sample_rate = sample_rate
+
+        perturbations = []
+        if speed_perturbation is not None:
+            perturbations.append(SpeedPerturbation(**speed_perturbation))
+        if gain_perturbation is not None:
+            perturbations.append(GainPerturbation(**gain_perturbation))
+        if shift_perturbation is not None:
+            perturbations.append(ShiftPerturbation(**shift_perturbation))
+        self.perturbations = perturbations
+
+        self.max_duration = max_duration
+
+        self.samples = []
+        self.duration = 0.0
+        self.duration_filtered = 0.0
+
+        for fpath in manifest_fpaths:
+            self._load_json_manifest(fpath)
+
+        if sort_by_duration:
+            self.samples = sorted(self.samples, key=lambda s: s['duration'])
+
+    def __getitem__(self, index):
+        s = self.samples[index]
+        rn_indx = np.random.randint(len(s['audio_filepath']))
+        duration = s['audio_duration'][rn_indx] if 'audio_duration' in s else 0
+        offset = s.get('offset', 0)
+
+        segment = AudioSegment(
+            s['audio_filepath'][rn_indx], target_sr=self.sample_rate,
+            offset=offset, duration=duration, trim=self.trim_silence)
+
+        for p in self.perturbations:
+            p.maybe_apply(segment, self.sample_rate)
+
+        segment = torch.FloatTensor(segment.samples)
+
+        return (segment,
+                torch.tensor(segment.shape[0]).int(),
+                torch.tensor(s["transcript"]),
+                torch.tensor(len(s["transcript"])).int())
+
+    def __len__(self):
+        return len(self.samples)
+
+    def _load_json_manifest(self, fpath):
+        for s in json.load(open(fpath, "r", encoding="utf-8")):
+
+            if self.pad_to_max_duration and not self.ignore_offline_speed_perturbation:
+                # require all perturbed samples to be < self.max_duration
+                s_max_duration = max(f['duration'] for f in s['files'])
+            else:
+                # otherwise we allow perturbances to be > self.max_duration
+                s_max_duration = s['original_duration']
+
+            s['duration'] = s.pop('original_duration')
+            if not (self.min_duration <= s_max_duration <= self.max_duration):
+                self.duration_filtered += s['duration']
+                continue
+
+            # Prune and normalize according to transcript
+            tr = (s.get('transcript', None) or
+                  self.load_transcript(s['text_filepath']))
+
+            if not isinstance(tr, str):
+                print(f'WARNING: Skipped sample (transcript not a str): {tr}.')
+                self.duration_filtered += s['duration']
+                continue
+
+            if self.normalize_transcripts:
+                tr = normalize_string(tr, self.labels, self.punctuation_map)
+
+            s["transcript"] = self.to_vocab_inds(tr)
+
+            files = s.pop('files')
+            if self.ignore_offline_speed_perturbation:
+                files = [f for f in files if f['speed'] == 1.0]
+
+            s['audio_duration'] = [f['duration'] for f in files]
+            s['audio_filepath'] = [str(Path(self.data_dir, f['fname']))
+                                   for f in files]
+            self.samples.append(s)
+            self.duration += s['duration']
+
+            if self.max_utts > 0 and len(self.samples) >= self.max_utts:
+                print(f'Reached max_utts={self.max_utts}. Finished parsing {fpath}.')
+                break
+
+    def load_transcript(self, transcript_path):
+        with open(transcript_path, 'r', encoding="utf-8") as transcript_file:
+            transcript = transcript_file.read().replace('\n', '')
+        return transcript
+
+    def to_vocab_inds(self, transcript):
+        chars = [self.labels_map.get(x, self.blank_index) for x in list(transcript)]
+        transcript = list(filter(lambda x: x != self.blank_index, chars))
+        return transcript
+
+
+def collate_fn(batch):
+    bs = len(batch)
+    max_len = lambda l, idx: max(el[idx].size(0) for el in l)
+
+    # audio = torch.zeros(bs, max_len(batch, 0))
+    audio = torch.zeros(bs,680000)
+    audio_lens = torch.zeros(bs, dtype=torch.int32)
+    # transcript = torch.zeros(bs, max_len(batch, 2))
+    transcript = torch.zeros(bs, 700)
+    transcript_lens = torch.zeros(bs, dtype=torch.int32)
+
+    for i, sample in enumerate(batch):
+        audio[i].narrow(0, 0, sample[0].size(0)).copy_(sample[0])
+        audio_lens[i] = sample[1]
+        transcript[i].narrow(0, 0, sample[2].size(0)).copy_(sample[2])
+        transcript_lens[i] = sample[3]
+    return audio, audio_lens, transcript, transcript_lens
+
+
+def get_data_loader(dataset, batch_size, multi_gpu=True, shuffle=True,
+                    drop_last=True, num_workers=4):
+
+    kw = {'dataset': dataset, 'collate_fn': collate_fn,
+          'num_workers': num_workers, 'pin_memory': True}
+
+    if multi_gpu:
+        loader_shuffle = False
+        sampler = DistributedSampler(dataset, shuffle=shuffle)
+    else:
+        loader_shuffle = shuffle
+        sampler = None
+
+    return DataLoader(batch_size=batch_size, drop_last=drop_last,
+                      sampler=sampler, shuffle=loader_shuffle, **kw)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/features.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/features.py
new file mode 100644
index 0000000000..731d765c24
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/features.py
@@ -0,0 +1,522 @@
+import math
+import random
+
+import librosa
+import torch
+import torch.nn as nn
+
+from apex import amp
+import numpy as np
+import torch.nn.functional as F
+from scipy.signal import get_window
+from librosa.util import pad_center, tiny, utils
+
+
+def window_sumsquare(window, n_frames, hop_length=200, win_length=800,
+                     n_fft=800, dtype=np.float32, norm=None):
+    """
+    # from librosa 0.6
+    Compute the sum-square envelope of a window function at a given hop length.
+    This is used to estimate modulation effects induced by windowing
+    observations in short-time fourier transforms.
+    Parameters
+    ----------
+    window : string, tuple, number, callable, or list-like
+        Window specification, as in `get_window`
+    n_frames : int > 0
+        The number of analysis frames
+    hop_length : int > 0
+        The number of samples to advance between frames
+    win_length : [optional]
+        The length of the window function.  By default, this matches `n_fft`.
+    n_fft : int > 0
+        The length of each analysis frame.
+    dtype : np.dtype
+        The data type of the output
+    Returns
+    -------
+    wss : np.ndarray, shape=`(n_fft + hop_length * (n_frames - 1))`
+        The sum-squared envelope of the window function
+    """
+    if win_length is None:
+        win_length = n_fft
+
+    n = n_fft + hop_length * (n_frames - 1)
+    x = np.zeros(n, dtype=dtype)
+
+    # Compute the squared window at the desired length
+    win_sq = get_window(window, win_length, fftbins=True)
+    win_sq = utils.normalize(win_sq, norm=norm) ** 2
+    win_sq = utils.pad_center(win_sq, n_fft)
+
+    # Fill the envelope
+    for i in range(n_frames):
+        sample = i * hop_length
+        x[sample:min(n, sample + n_fft)] += win_sq[:max(0, min(n_fft, n - sample))]
+    return x
+
+
+class STFT(torch.nn.Module):
+    def __init__(self, filter_length=1024, hop_length=512, win_length=None,
+                 window='hann'):
+        """
+        This module implements an STFT using 1D convolution and 1D transpose convolutions.
+        This is a bit tricky so there are some cases that probably won't work as working
+        out the same sizes before and after in all overlap add setups is tough. Right now,
+        this code should work with hop lengths that are half the filter length (50% overlap
+        between frames).
+
+        Keyword Arguments:
+            filter_length {int} -- Length of filters used (default: {1024})
+            hop_length {int} -- Hop length of STFT (restrict to 50% overlap between frames) (default: {512})
+            win_length {[type]} -- Length of the window function applied to each frame (if not specified, it
+                equals the filter length). (default: {None})
+            window {str} -- Type of window to use (options are bartlett, hann, hamming, blackman, blackmanharris)
+                (default: {'hann'})
+        """
+        super(STFT, self).__init__()
+        self.filter_length = filter_length
+        self.hop_length = hop_length
+        self.win_length = win_length if win_length else filter_length
+        self.window = window
+        self.forward_transform = None
+        self.pad_amount = int(self.filter_length / 2)
+        scale = self.filter_length / self.hop_length
+        fourier_basis = np.fft.fft(np.eye(self.filter_length))
+
+        cutoff = int((self.filter_length / 2 + 1))
+        fourier_basis = np.vstack([np.real(fourier_basis[:cutoff, :]),
+                                   np.imag(fourier_basis[:cutoff, :])])
+        forward_basis = torch.FloatTensor(fourier_basis[:, None, :])
+        inverse_basis = torch.FloatTensor(
+            np.linalg.pinv(scale * fourier_basis).T[:, None, :])
+
+        assert (filter_length >= self.win_length)
+        # get window and zero center pad it to filter_length
+        fft_window = get_window(window, self.win_length, fftbins=True)
+        fft_window = pad_center(fft_window, filter_length)
+        fft_window = torch.from_numpy(fft_window).float()
+
+        # window the bases
+        forward_basis *= fft_window
+        inverse_basis *= fft_window
+
+        self.register_buffer('forward_basis', forward_basis.float())
+        self.register_buffer('inverse_basis', inverse_basis.float())
+
+    def transform(self, input_data):
+        """Take input data (audio) to STFT domain.
+
+        Arguments:
+            input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
+
+        Returns:
+            magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+            phase {tensor} -- Phase of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+        """
+        num_batches = input_data.shape[0]
+        num_samples = input_data.shape[-1]
+
+        self.num_samples = num_samples
+
+        # similar to librosa, reflect-pad the input
+        input_data = input_data.view(num_batches, 1, num_samples)
+
+        input_data = F.pad(
+            input_data.unsqueeze(1),
+            (self.pad_amount, self.pad_amount, 0, 0),
+            mode='constant')
+        input_data = input_data.squeeze(1)
+
+        forward_transform = F.conv1d(
+            input_data,
+            self.forward_basis,
+            stride=self.hop_length,
+            padding=0)
+
+        cutoff = int((self.filter_length / 2) + 1)
+        real_part = forward_transform[:, :cutoff, :]
+        imag_part = forward_transform[:, cutoff:, :]
+
+        magnitude = torch.sqrt(real_part ** 2 + imag_part ** 2)
+        phase = torch.atan2(imag_part.data, real_part.data)
+
+        return magnitude, phase
+
+    def inverse(self, magnitude, phase):
+        """Call the inverse STFT (iSTFT), given magnitude and phase tensors produced
+        by the ```transform``` function.
+
+        Arguments:
+            magnitude {tensor} -- Magnitude of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+            phase {tensor} -- Phase of STFT with shape (num_batch,
+                num_frequencies, num_frames)
+
+        Returns:
+            inverse_transform {tensor} -- Reconstructed audio given magnitude and phase. Of
+                shape (num_batch, num_samples)
+        """
+        recombine_magnitude_phase = torch.cat(
+            [magnitude * torch.cos(phase), magnitude * torch.sin(phase)], dim=1)
+
+        inverse_transform = F.conv_transpose1d(
+            recombine_magnitude_phase,
+            self.inverse_basis,
+            stride=self.hop_length,
+            padding=0)
+
+        if self.window is not None:
+            window_sum = window_sumsquare(
+                self.window, magnitude.size(-1), hop_length=self.hop_length,
+                win_length=self.win_length, n_fft=self.filter_length,
+                dtype=np.float32)
+            # remove modulation effects
+            approx_nonzero_indices = torch.from_numpy(
+                np.where(window_sum > tiny(window_sum))[0])
+            window_sum = torch.from_numpy(window_sum).to(inverse_transform.device)
+            inverse_transform[:, :, approx_nonzero_indices] /= window_sum[approx_nonzero_indices]
+
+            # scale by hop ratio
+            inverse_transform *= float(self.filter_length) / self.hop_length
+
+        inverse_transform = inverse_transform[..., self.pad_amount:]
+        inverse_transform = inverse_transform[..., :self.num_samples]
+        inverse_transform = inverse_transform.squeeze(1)
+
+        return inverse_transform
+
+    def forward(self, input_data):
+        """Take input data (audio) to STFT domain and then back to audio.
+
+        Arguments:
+            input_data {tensor} -- Tensor of floats, with shape (num_batch, num_samples)
+
+        Returns:
+            reconstruction {tensor} -- Reconstructed audio given magnitude and phase. Of
+                shape (num_batch, num_samples)
+        """
+        # print("input_data",input_data)
+        self.magnitude, self.phase = self.transform(input_data)
+        reconstruction = self.inverse(self.magnitude, self.phase)
+        return reconstruction
+# stft = STFT()
+class BaseFeatures(nn.Module):
+    """Base class for GPU accelerated audio preprocessing."""
+    __constants__ = ["pad_align", "pad_to_max_duration", "max_len"]
+
+    def __init__(self, pad_align, pad_to_max_duration, max_duration,
+                 sample_rate, window_size, window_stride, spec_augment=None,
+                 cutout_augment=None):
+        super(BaseFeatures, self).__init__()
+
+        self.pad_align = pad_align
+        self.pad_to_max_duration = pad_to_max_duration
+        self.win_length = int(sample_rate * window_size) # frame size
+        self.hop_length = int(sample_rate * window_stride)
+
+        # Calculate maximum sequence length (# frames)
+        if pad_to_max_duration:
+            self.max_len = 1 + math.ceil(
+                (max_duration * sample_rate - self.win_length) / self.hop_length
+            )
+
+        if spec_augment is not None:
+            self.spec_augment = SpecAugment(**spec_augment)
+        else:
+            self.spec_augment = None
+
+        if cutout_augment is not None:
+            self.cutout_augment = CutoutAugment(**cutout_augment)
+        else:
+            self.cutout_augment = None
+
+    @torch.no_grad()
+    def calculate_features(self, audio, audio_lens):
+        return audio, audio_lens
+
+    def __call__(self, audio, audio_lens, optim_level=0):
+        dtype = audio.dtype
+        audio = audio.float()
+        if optim_level == 1:
+            with amp.disable_casts():
+                feat, feat_lens = self.calculate_features(audio, audio_lens)
+        else:
+            feat, feat_lens = self.calculate_features(audio, audio_lens)
+
+        feat = self.apply_padding(feat)
+
+        if self.cutout_augment is not None:
+            feat = self.cutout_augment(feat)
+
+        if self.spec_augment is not None:
+            feat = self.spec_augment(feat)
+
+        feat = feat.to(dtype)
+        return feat, feat_lens
+
+    def apply_padding(self, x):
+        if self.pad_to_max_duration:
+            x_size = max(x.size(-1), self.max_len)
+        else:
+            x_size = x.size(-1)
+
+        if self.pad_align > 0:
+            pad_amt = x_size % self.pad_align
+        else:
+            pad_amt = 0
+
+        padded_len = x_size + (self.pad_align - pad_amt if pad_amt > 0 else 0)
+        return nn.functional.pad(x, (0, padded_len - x.size(-1)))
+
+class SpecAugment(nn.Module):
+    """Spec augment. refer to https://arxiv.org/abs/1904.08779
+    """
+    def __init__(self, freq_masks=0, min_freq=0, max_freq=10, time_masks=0,
+                 min_time=0, max_time=10):
+        super(SpecAugment, self).__init__()
+        assert 0 <= min_freq <= max_freq
+        assert 0 <= min_time <= max_time
+
+        self.freq_masks = freq_masks
+        self.min_freq = min_freq
+        self.max_freq = max_freq
+
+        self.time_masks = time_masks
+        self.min_time = min_time
+        self.max_time = max_time
+
+    @torch.no_grad()
+    def forward(self, x):
+        sh = x.shape
+        mask = torch.zeros(x.shape, dtype=torch.bool, device=x.device)
+
+        for idx in range(sh[0]):
+            for _ in range(self.freq_masks):
+                w = torch.randint(self.min_freq, self.max_freq + 1, size=(1,)).item()
+                f0 = torch.randint(0, max(1, sh[1] - w), size=(1,))
+                mask[idx, f0:f0+w] = 1
+
+            for _ in range(self.time_masks):
+                w = torch.randint(self.min_time, self.max_time + 1, size=(1,)).item()
+                t0 = torch.randint(0, max(1, sh[2] - w), size=(1,))
+                mask[idx, :, t0:t0+w] = 1
+
+        return x.masked_fill(mask, 0)
+
+
+class CutoutAugment(nn.Module):
+    """Cutout. refer to https://arxiv.org/pdf/1708.04552.pdf
+    """
+    def __init__(self, masks=0, min_freq=20, max_freq=20, min_time=5, max_time=5):
+        super(CutoutAugment, self).__init__()
+        assert 0 <= min_freq <= max_freq
+        assert 0 <= min_time <= max_time
+
+        self.masks = masks
+        self.min_freq = min_freq
+        self.max_freq = max_freq
+        self.min_time = min_time
+        self.max_time = max_time
+
+    @torch.no_grad()
+    def forward(self, x):
+        sh = x.shape
+        mask = torch.zeros(x.shape, dtype=torch.bool, device=x.device)
+
+        for idx in range(sh[0]):
+            for i in range(self.masks):
+
+                w = torch.randint(self.min_freq, self.max_freq + 1, size=(1,)).item()
+                h = torch.randint(self.min_time, self.max_time + 1, size=(1,)).item()
+
+                f0 = int(random.uniform(0, sh[1] - w))
+                t0 = int(random.uniform(0, sh[2] - h))
+
+                mask[idx, f0:f0+w, t0:t0+h] = 1
+
+        return x.masked_fill(mask, 0)
+
+
+@torch.jit.script
+def normalize_batch(x, seq_len, normalize_type: str):
+#    print ("normalize_batch: x, seq_len, shapes: ", x.shape, seq_len, seq_len.shape)
+    if normalize_type == "per_feature":
+        x_mean = torch.zeros((seq_len.shape[0], x.shape[1]), dtype=x.dtype,
+                                                 device=x.device)
+        x_std = torch.zeros((seq_len.shape[0], x.shape[1]), dtype=x.dtype,
+                                                device=x.device)
+        for i in range(x.shape[0]):
+            x_mean[i, :] = x[i, :, :seq_len[i]].mean(dim=1)
+            x_std[i, :] = x[i, :, :seq_len[i]].std(dim=1)
+        # make sure x_std is not zero
+        x_std += 1e-5
+        return (x - x_mean.unsqueeze(2)) / x_std.unsqueeze(2)
+
+    elif normalize_type == "all_features":
+        x_mean = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device)
+        x_std = torch.zeros(seq_len.shape, dtype=x.dtype, device=x.device)
+        for i in range(x.shape[0]):
+            x_mean[i] = x[i, :, :int(seq_len[i])].mean()
+            x_std[i] = x[i, :, :int(seq_len[i])].std()
+        # make sure x_std is not zero
+        x_std += 1e-5
+        return (x - x_mean.view(-1, 1, 1)) / x_std.view(-1, 1, 1)
+    else:
+        return x
+
+
+@torch.jit.script
+def stack_subsample_frames(x, x_lens, stacking: int = 1, subsampling: int = 1):
+    """ Stacks frames together across feature dim, and then subsamples
+
+    input is batch_size, feature_dim, num_frames
+    output is batch_size, feature_dim * stacking, num_frames / subsampling
+
+    """
+    seq = [x]
+    for n in range(1, stacking):
+        tmp = torch.zeros_like(x)
+        tmp[:, :, :-n] = x[:, :, n:]
+        seq.append(tmp)
+    x = torch.cat(seq, dim=1)[:, :, ::subsampling]
+
+    if subsampling > 1:
+        x_lens = torch.ceil(x_lens.float() / subsampling).int()
+
+        if x.size(2) > x_lens.max().item():
+            assert abs(x.size(2) - x_lens.max().item()) <= 1
+            x = x[:,:,:x_lens.max().item()]
+
+    return x, x_lens
+
+
+class FilterbankFeatures(BaseFeatures):
+    # For JIT, https://pytorch.org/docs/stable/jit.html#python-defined-constants
+    __constants__ = ["dither", "preemph", "n_fft", "hop_length", "win_length",
+                     "log", "frame_splicing", "normalize"]
+    # torchscript: "center" removed due to a bug
+
+    def __init__(self, spec_augment=None, cutout_augment=None,
+                 sample_rate=8000, window_size=0.02, window_stride=0.01,
+                 window="hamming", normalize="per_feature", n_fft=None,
+                 preemph=0.97, n_filt=64, lowfreq=0, highfreq=None, log=True,
+                 dither=1e-5, pad_align=8, pad_to_max_duration=False,
+                 max_duration=float('inf'), frame_splicing=1):
+        super(FilterbankFeatures, self).__init__(
+            pad_align=pad_align, pad_to_max_duration=pad_to_max_duration,
+            max_duration=max_duration, sample_rate=sample_rate,
+            window_size=window_size, window_stride=window_stride,
+            spec_augment=spec_augment, cutout_augment=cutout_augment)
+
+        torch_windows = {
+            'hann': torch.hann_window,
+            'hamming': torch.hamming_window,
+            'blackman': torch.blackman_window,
+            'bartlett': torch.bartlett_window,
+            'none': None,
+        }
+
+        self.n_fft = n_fft or 2 ** math.ceil(math.log2(self.win_length))
+
+        self.normalize = normalize
+        self.log = log
+        #TORCHSCRIPT: Check whether or not we need this
+        self.dither = dither
+        self.frame_splicing = frame_splicing
+        self.n_filt = n_filt
+        self.preemph = preemph
+        highfreq = highfreq or sample_rate / 2
+        window_fn = torch_windows.get(window, None)
+        window_tensor = window_fn(self.win_length,
+                                  periodic=False) if window_fn else None
+        filterbanks = torch.tensor(
+            librosa.filters.mel(sample_rate, self.n_fft, n_mels=n_filt,
+                                fmin=lowfreq, fmax=highfreq),
+            dtype=torch.float).unsqueeze(0)
+        # torchscript
+        self.register_buffer("fb", filterbanks)
+        self.register_buffer("window", window_tensor)
+        # print("n_fft", self.n_fft)
+        # print("hop_length", self.hop_length)
+        # print("win_length", self.win_length)
+        # print("window", self.window.to(dtype=torch.float))
+
+        # self.stft = STFT(filter_length=512, hop_length=160,win_length=320)
+    def get_seq_len(self, seq_len):
+        return torch.ceil(seq_len.to(dtype=torch.float) / self.hop_length).to(
+            dtype=torch.int)
+
+    # do stft
+    # TORCHSCRIPT: center removed due to bug
+    # def stft(self, x):
+    #     return torch.stft(x, n_fft=self.n_fft, hop_length=self.hop_length,
+    #                       win_length=self.win_length,pad_mode = "constant",
+    #                       window=self.window.to(dtype=torch.float))
+
+    def stft(self, x):
+        result = []
+        for i in range(x.shape[0]):
+            tmp = librosa.stft(x[i].numpy(), n_fft=self.n_fft, hop_length=self.hop_length,
+                               win_length=self.win_length, pad_mode="reflect",
+                               window='hann')
+            tmp_real = torch.from_numpy(tmp.real).float()
+            tmp_imag = torch.from_numpy(tmp.imag).float()
+            tmp = torch.stack((tmp_real, tmp_imag), dim=2)
+            result.append(tmp)
+
+        return torch.stack((result), dim=0)
+    @torch.no_grad()
+    def calculate_features(self, x, seq_len):
+        dtype = x.dtype
+        seq_len = self.get_seq_len(seq_len)
+        # dither
+        # print(seq_len)
+        if self.dither > 0:
+            x += self.dither * torch.randn_like(x)
+
+        # do preemphasis
+        if self.preemph is not None:
+            x = torch.cat(
+                (x[:, 0].unsqueeze(1), x[:, 1:] - self.preemph * x[:, :-1]), dim=1)
+        # print(x)
+        # print(seq_len)
+        # x = x.to("cpu")
+        # print("inputxsize",x.size())
+        # x = x.numpy()
+        x  = self.stft(x)
+        # x = torch.from_numpy(x)
+        # x = x.npu()
+            # get power spectrum
+        x = x.pow(2).sum(-1)
+        # print("x-stft-size:", x.size())
+        # print(x.dtype)
+        y = self.fb.to(x.dtype)
+        # print("fb_ysize:", y.size())
+        # dot with filterbank energies
+        x = torch.matmul(self.fb.to(x.dtype), x)
+
+        # log features if required
+        if self.log:
+            x = torch.log(x + 1e-20)
+
+        # frame splicing if required
+        if self.frame_splicing > 1:
+            raise ValueError('Frame splicing not supported')
+
+        # normalize if required
+        x = normalize_batch(x, seq_len, normalize_type=self.normalize)
+
+        # mask to zero any values beyond seq_len in batch,
+        # pad to multiple of `pad_align` (for efficiency)
+        max_len = x.size(-1)
+        mask = torch.arange(max_len, dtype=seq_len.dtype, device=x.device)
+        mask = mask.expand(x.size(0), max_len) >= seq_len.unsqueeze(1)
+        x = x.masked_fill(mask.unsqueeze(1), 0)
+
+        # TORCHSCRIPT: Is this del important? It breaks scripting
+        # del mask
+
+        return x.to(dtype), seq_len
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/helpers.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/helpers.py
new file mode 100644
index 0000000000..efdb11a445
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/helpers.py
@@ -0,0 +1,300 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import glob
+import os
+import re
+from collections import OrderedDict
+
+from apex import amp
+
+import torch
+import torch.distributed as dist
+
+from .metrics import word_error_rate
+
+
+def print_once(msg):
+    if not dist.is_initialized() or dist.get_rank() == 0:
+        print(msg)
+
+
+def add_ctc_blank(symbols):
+    return symbols + ['<BLANK>']
+
+
+def ctc_decoder_predictions_tensor(tensor, labels):
+    """
+    Takes output of greedy ctc decoder and performs ctc decoding algorithm to
+    remove duplicates and special symbol. Returns prediction
+    Args:
+        tensor: model output tensor
+        label: A list of labels
+    Returns:
+        prediction
+    """
+    blank_id = len(labels) - 1
+    hypotheses = []
+    labels_map = {i: labels[i] for i in range(len(labels))}
+    prediction_cpu_tensor = tensor.long().cpu()
+    # iterate over batch
+    for ind in range(prediction_cpu_tensor.shape[0]):
+        prediction = prediction_cpu_tensor[ind].numpy().tolist()
+        # CTC decoding procedure
+        decoded_prediction = []
+        previous = len(labels) - 1 # id of a blank symbol
+        for p in prediction:
+            if (p != previous or previous == blank_id) and p != blank_id:
+                decoded_prediction.append(p)
+            previous = p
+        hypothesis = ''.join([labels_map[c] for c in decoded_prediction])
+        hypotheses.append(hypothesis)
+    return hypotheses
+
+
+def greedy_wer(preds, tgt, tgt_lens, labels):
+    """
+    Takes output of greedy ctc decoder and performs ctc decoding algorithm to
+    remove duplicates and special symbol. Prints wer and prediction examples to screen
+    Args:
+        tensors: A list of 3 tensors (predictions, targets, target_lengths)
+        labels: A list of labels
+
+    Returns:
+        word error rate
+    """
+    with torch.no_grad():
+        references = gather_transcripts([tgt], [tgt_lens], labels)
+        hypotheses = ctc_decoder_predictions_tensor(preds, labels)
+
+    wer, _, _ = word_error_rate(hypotheses, references)
+    return wer, hypotheses[0], references[0]
+
+
+def gather_losses(losses_list):
+    return [torch.mean(torch.stack(losses_list))]
+
+
+def gather_predictions(predictions_list, labels):
+    results = []
+    for prediction in predictions_list:
+        results += ctc_decoder_predictions_tensor(prediction, labels=labels)
+    return results
+
+
+def gather_transcripts(transcript_list, transcript_len_list, labels):
+    results = []
+    labels_map = {i: labels[i] for i in range(len(labels))}
+    # iterate over workers
+    for txt, lens in zip(transcript_list, transcript_len_list):
+        for t, l in zip(txt.long().cpu(), lens.long().cpu()):
+            t = list(t.numpy())
+            results.append(''.join([labels_map[c] for c in t[:l]]))
+    return results
+
+
+def process_evaluation_batch(tensors, global_vars, labels):
+    """
+    Processes results of an iteration and saves it in global_vars
+    Args:
+        tensors: dictionary with results of an evaluation iteration, e.g. loss, predictions, transcript, and output
+        global_vars: dictionary where processes results of iteration are saved
+        labels: A list of labels
+    """
+    for kv, v in tensors.items():
+        if kv.startswith('loss'):
+            global_vars['EvalLoss'] += gather_losses(v)
+        elif kv.startswith('predictions'):
+            global_vars['preds'] += gather_predictions(v, labels)
+        elif kv.startswith('transcript_length'):
+            transcript_len_list = v
+        elif kv.startswith('transcript'):
+            transcript_list = v
+        elif kv.startswith('output'):
+            global_vars['logits'] += v
+
+    global_vars['txts'] += gather_transcripts(
+        transcript_list, transcript_len_list, labels)
+
+
+def process_evaluation_epoch(aggregates, tag=None):
+    """
+    Processes results from each worker at the end of evaluation and combine to final result
+    Args:
+        aggregates: dictionary containing information of entire evaluation
+    Return:
+        wer: final word error rate
+        loss: final loss
+    """
+    if 'losses' in aggregates:
+        eloss = torch.mean(torch.stack(aggregates['losses'])).item()
+    else:
+        eloss = None
+    hypotheses = aggregates['preds']
+    references = aggregates['txts']
+
+    wer, scores, num_words = word_error_rate(hypotheses, references)
+    multi_gpu = dist.is_initialized()
+    if multi_gpu:
+        if eloss is not None:
+            eloss /= dist.get_world_size()
+            eloss_tensor = torch.tensor(eloss).npu()
+            dist.all_reduce(eloss_tensor)
+            eloss = eloss_tensor.item()
+
+        scores_tensor = torch.tensor(scores).npu().float()
+        dist.all_reduce(scores_tensor)
+        scores = scores_tensor.item()
+        num_words_tensor = torch.tensor(num_words).npu().float()
+        dist.all_reduce(num_words_tensor)
+        num_words = num_words_tensor.item()
+        wer = scores * 1.0 / num_words
+    return wer, eloss
+
+
+def num_weights(module):
+    return sum(p.numel() for p in module.parameters() if p.requires_grad)
+
+
+def convert_v1_state_dict(state_dict):
+    rules = [
+        ('^jasper_encoder.encoder.', 'encoder.layers.'),
+        ('^jasper_decoder.decoder_layers.', 'decoder.layers.'),
+    ]
+    ret = {}
+    for k, v in state_dict.items():
+        if k.startswith('acoustic_model.'):
+            continue
+        if k.startswith('audio_preprocessor.'):
+            continue
+        for pattern, to in rules:
+            k = re.sub(pattern, to, k)
+        ret[k] = v
+
+    return ret
+
+
+class Checkpointer(object):
+
+    def __init__(self, save_dir, model_name, keep_milestones=[100,200,300],
+                 use_amp=False):
+        self.save_dir = save_dir
+        self.keep_milestones = keep_milestones
+        self.use_amp = use_amp
+        self.model_name = model_name
+
+        tracked = [
+            (int(re.search('epoch(\d+)_', f).group(1)), f)
+            for f in glob.glob(f'{save_dir}/{self.model_name}_epoch*_checkpoint.pt')]
+        tracked = sorted(tracked, key=lambda t: t[0])
+        self.tracked = OrderedDict(tracked)
+
+    def save(self, model, ema_model, optimizer, epoch, step, best_wer,
+             is_best=False):
+        """Saves model checkpoint for inference/resuming training.
+
+        Args:
+            model: the model, optionally wrapped by DistributedDataParallel
+            ema_model: model with averaged weights, can be None
+            optimizer: optimizer
+            epoch (int): epoch during which the model is saved
+            step (int): number of steps since beginning of training
+            best_wer (float): lowest recorded WER on the dev set
+            is_best (bool, optional): set name of checkpoint to 'best'
+                and overwrite the previous one
+        """
+        rank = 0
+        if dist.is_initialized():
+            dist.barrier()
+            rank = dist.get_rank()
+
+        if rank != 0:
+            return
+
+        # Checkpoint already saved
+        if not is_best and epoch in self.tracked:
+            return
+
+        unwrap_ddp = lambda model: getattr(model, 'module', model)
+        state = {
+            'epoch': epoch,
+            'step': step,
+            'best_wer': best_wer,
+            'state_dict': unwrap_ddp(model).state_dict(),
+            'ema_state_dict': unwrap_ddp(ema_model).state_dict() if ema_model is not None else None,
+            'optimizer': optimizer.state_dict(),
+            'amp': amp.state_dict() if self.use_amp else None,
+        }
+
+        if is_best:
+            fpath = os.path.join(
+                self.save_dir, f"{self.model_name}_best_checkpoint.pt")
+        else:
+            fpath = os.path.join(
+                self.save_dir, f"{self.model_name}_epoch{epoch}_checkpoint.pt")
+
+        print_once(f"Saving {fpath}...")
+        torch.save(state, fpath)
+
+        if not is_best:
+            # Remove old checkpoints; keep milestones and the last two
+            self.tracked[epoch] = fpath
+            for epoch in set(list(self.tracked)[:-2]) - set(self.keep_milestones):
+                try:
+                    os.remove(self.tracked[epoch])
+                except:
+                    pass
+                del self.tracked[epoch]
+
+    def last_checkpoint(self):
+        tracked = list(self.tracked.values())
+
+        if len(tracked) >= 1:
+            try:
+                torch.load(tracked[-1], map_location='cpu')
+                return tracked[-1]
+            except:
+                print_once(f'Last checkpoint {tracked[-1]} appears corrupted.')
+
+        elif len(tracked) >= 2:
+            return tracked[-2]
+        else:
+            return None
+
+    def load(self, fpath, model, ema_model, optimizer, meta):
+
+        print_once(f'Loading model from {fpath}')
+        checkpoint = torch.load(fpath, map_location="cpu")
+
+        unwrap_ddp = lambda model: getattr(model, 'module', model)
+        state_dict = convert_v1_state_dict(checkpoint['state_dict'])
+        unwrap_ddp(model).load_state_dict(state_dict, strict=True)
+
+        if ema_model is not None:
+            if checkpoint.get('ema_state_dict') is not None:
+                key = 'ema_state_dict'
+            else:
+                key = 'state_dict'
+                print_once('WARNING: EMA weights not found in the checkpoint.')
+                print_once('WARNING: Initializing EMA model with regular params.')
+            state_dict = convert_v1_state_dict(checkpoint[key])
+            unwrap_ddp(ema_model).load_state_dict(state_dict, strict=True)
+
+        optimizer.load_state_dict(checkpoint['optimizer'])
+
+        if self.use_amp:
+            amp.load_state_dict(checkpoint['amp'])
+
+        meta['start_epoch'] = checkpoint.get('epoch')
+        meta['best_wer'] = checkpoint.get('best_wer', meta['best_wer'])
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/metrics.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/metrics.py
new file mode 100644
index 0000000000..4ae47a4c06
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/metrics.py
@@ -0,0 +1,59 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+
+def __levenshtein(a, b):
+    """Calculates the Levenshtein distance between two sequences."""
+
+    n, m = len(a), len(b)
+    if n > m:
+        # Make sure n <= m, to use O(min(n,m)) space
+        a, b = b, a
+        n, m = m, n
+
+    current = list(range(n + 1))
+    for i in range(1, m + 1):
+        previous, current = current, [i] + [0] * n
+        for j in range(1, n + 1):
+            add, delete = previous[j] + 1, current[j - 1] + 1
+            change = previous[j - 1]
+            if a[j - 1] != b[i - 1]:
+                change = change + 1
+            current[j] = min(add, delete, change)
+
+    return current[n]
+
+
+def word_error_rate(hypotheses, references):
+    """Computes average Word Error Rate (WER) between two text lists."""
+
+    scores = 0
+    words = 0
+    len_diff = len(references) - len(hypotheses)
+    if len_diff > 0:
+        raise ValueError("Uneqal number of hypthoses and references: "
+                         "{0} and {1}".format(len(hypotheses), len(references)))
+    elif len_diff < 0:
+        hypotheses = hypotheses[:len_diff]
+
+    for h, r in zip(hypotheses, references):
+        h_list = h.split()
+        r_list = r.split()
+        words += len(r_list)
+        scores += __levenshtein(h_list, r_list)
+    if words!=0:
+        wer = 1.0*scores/words
+    else:
+        wer = float('inf')
+    return wer, scores, words
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/optimizers.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/optimizers.py
new file mode 100644
index 0000000000..8175919196
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/optimizers.py
@@ -0,0 +1,269 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import torch
+from torch.optim import Optimizer
+import math
+
+
+def lr_policy(step, epoch, initial_lr, optimizer, steps_per_epoch, warmup_epochs,
+              hold_epochs, num_epochs=None, policy='linear', min_lr=1e-5,
+              exp_gamma=None):
+    """
+    learning rate decay
+    Args:
+        initial_lr: base learning rate
+        step: current iteration number
+        N: total number of iterations over which learning rate is decayed
+        lr_steps: list of steps to apply exp_gamma
+    """
+    warmup_steps = warmup_epochs * steps_per_epoch
+    hold_steps = hold_epochs * steps_per_epoch
+
+    if policy == 'legacy':
+        assert num_epochs is not None
+        tot_steps = num_epochs * steps_per_epoch
+
+        if step < warmup_steps:
+            a = (step + 1) / (warmup_steps + 1)
+        elif step < warmup_steps + hold_steps:
+            a = 1.0
+        else:
+            a = (((tot_steps - step)
+                 / (tot_steps - warmup_steps - hold_steps)) ** 2)
+
+    elif policy == 'exponential':
+        assert exp_gamma is not None
+
+        if step < warmup_steps:
+            a = (step + 1) / (warmup_steps + 1)
+        elif step < warmup_steps + hold_steps:
+            a = 1.0
+        else:
+            a = exp_gamma ** (epoch - warmup_epochs - hold_epochs)
+
+    else:
+        raise ValueError
+
+    new_lr = max(a * initial_lr, min_lr)
+    for param_group in optimizer.param_groups:
+        param_group['lr'] = new_lr
+
+
+class AdamW(Optimizer):
+    """Implements AdamW algorithm.
+  
+    It has been proposed in `Adam: A Method for Stochastic Optimization`_.
+  
+    Arguments:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.9, 0.999))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+  
+        Adam: A Method for Stochastic Optimization:
+        https://arxiv.org/abs/1412.6980
+        On the Convergence of Adam and Beyond:
+        https://openreview.net/forum?id=ryQu7f-RZ
+    """
+  
+    def __init__(self, params, lr=1e-3, betas=(0.9, 0.999), eps=1e-8,
+                  weight_decay=0, amsgrad=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                        weight_decay=weight_decay, amsgrad=amsgrad)
+        super(AdamW, self).__init__(params, defaults)
+  
+    def __setstate__(self, state):
+        super(AdamW, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+  
+    def step(self, closure=None):
+        """Performs a single optimization step.
+  
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+                and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+  
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Adam does not support sparse gradients, please consider SparseAdam instead')
+                amsgrad = group['amsgrad']
+  
+                state = self.state[p]
+  
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros_like(p.data)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros_like(p.data)
+  
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                if amsgrad:
+                    max_exp_avg_sq = state['max_exp_avg_sq']
+                beta1, beta2 = group['betas']
+  
+                state['step'] += 1
+                # Decay the first and second moment running average coefficient
+                exp_avg.mul_(beta1).add_(grad, alpha=1 - beta1)
+                exp_avg_sq.mul_(beta2).addcmul_(1 - beta2, grad, grad)
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group['eps'])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+  
+                bias_correction1 = 1 - beta1 ** state['step']
+                bias_correction2 = 1 - beta2 ** state['step']
+                step_size = group['lr'] * math.sqrt(bias_correction2) / bias_correction1
+                p.data.add_(torch.mul(p.data, group['weight_decay']).addcdiv_(1, exp_avg, denom), alpha=-step_size)
+  
+        return loss
+
+  
+class Novograd(Optimizer):
+    """
+    Implements Novograd algorithm.
+
+    Args:
+        params (iterable): iterable of parameters to optimize or dicts defining
+            parameter groups
+        lr (float, optional): learning rate (default: 1e-3)
+        betas (Tuple[float, float], optional): coefficients used for computing
+            running averages of gradient and its square (default: (0.95, 0))
+        eps (float, optional): term added to the denominator to improve
+            numerical stability (default: 1e-8)
+        weight_decay (float, optional): weight decay (L2 penalty) (default: 0)
+        grad_averaging: gradient averaging
+        amsgrad (boolean, optional): whether to use the AMSGrad variant of this
+            algorithm from the paper `On the Convergence of Adam and Beyond`_
+            (default: False)
+    """
+
+    def __init__(self, params, lr=1e-3, betas=(0.95, 0), eps=1e-8,
+                 weight_decay=0, grad_averaging=False, amsgrad=False):
+        if not 0.0 <= lr:
+            raise ValueError("Invalid learning rate: {}".format(lr))
+        if not 0.0 <= eps:
+            raise ValueError("Invalid epsilon value: {}".format(eps))
+        if not 0.0 <= betas[0] < 1.0:
+            raise ValueError("Invalid beta parameter at index 0: {}".format(betas[0]))
+        if not 0.0 <= betas[1] < 1.0:
+            raise ValueError("Invalid beta parameter at index 1: {}".format(betas[1]))
+        defaults = dict(lr=lr, betas=betas, eps=eps,
+                      weight_decay=weight_decay,
+                      grad_averaging=grad_averaging,
+                      amsgrad=amsgrad)
+
+        super(Novograd, self).__init__(params, defaults)
+
+    def __setstate__(self, state):
+        super(Novograd, self).__setstate__(state)
+        for group in self.param_groups:
+            group.setdefault('amsgrad', False)
+
+    def step(self, closure=None):
+        """Performs a single optimization step.
+
+        Arguments:
+            closure (callable, optional): A closure that reevaluates the model
+            and returns the loss.
+        """
+        loss = None
+        if closure is not None:
+            loss = closure()
+
+        for group in self.param_groups:
+            for p in group['params']:
+                if p.grad is None:
+                    continue
+                grad = p.grad.data
+                if grad.is_sparse:
+                    raise RuntimeError('Sparse gradients are not supported.')
+                amsgrad = group['amsgrad']
+
+                state = self.state[p]
+
+                # State initialization
+                if len(state) == 0:
+                    state['step'] = 0
+                    # Exponential moving average of gradient values
+                    state['exp_avg'] = torch.zeros_like(p.data)
+                    # Exponential moving average of squared gradient values
+                    state['exp_avg_sq'] = torch.zeros([]).to(state['exp_avg'].device)
+                    if amsgrad:
+                        # Maintains max of all exp. moving avg. of sq. grad. values
+                        state['max_exp_avg_sq'] = torch.zeros([]).to(state['exp_avg'].device)
+
+                exp_avg, exp_avg_sq = state['exp_avg'], state['exp_avg_sq']
+                if amsgrad:
+                    max_exp_avg_sq = state['max_exp_avg_sq']
+                beta1, beta2 = group['betas']
+
+                state['step'] += 1
+
+                norm = torch.sum(torch.pow(grad, 2))
+
+                if exp_avg_sq == 0:
+                    exp_avg_sq.copy_(norm)
+                else:
+                    exp_avg_sq.mul_(beta2).add_(norm, alpha=1 - beta2)
+
+                if amsgrad:
+                    # Maintains the maximum of all 2nd moment running avg. till now
+                    torch.max(max_exp_avg_sq, exp_avg_sq, out=max_exp_avg_sq)
+                    # Use the max. for normalizing running avg. of gradient
+                    denom = max_exp_avg_sq.sqrt().add_(group['eps'])
+                else:
+                    denom = exp_avg_sq.sqrt().add_(group['eps'])
+
+                grad.div_(denom)
+                if group['weight_decay'] != 0:
+                    grad.add_(p.data, alpha=group['weight_decay'])
+                if group['grad_averaging']:
+                    grad.mul_(1 - beta1)
+                exp_avg.mul_(beta1).add_(grad)
+
+                p.data.add_(exp_avg, alpha=-group['lr'])
+        
+        return loss
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/tb_dllogger.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/tb_dllogger.py
new file mode 100644
index 0000000000..ecc6ec8689
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/tb_dllogger.py
@@ -0,0 +1,159 @@
+import atexit
+import glob
+import os
+import re
+import numpy as np
+
+import torch
+from torch.utils.tensorboard import SummaryWriter
+
+import dllogger
+from dllogger import StdOutBackend, JSONStreamBackend, Verbosity
+
+
+tb_loggers = {}
+
+
+class TBLogger:
+    """
+    xyz_dummies: stretch the screen with empty plots so the legend would
+                 always fit for other plots
+    """
+    def __init__(self, enabled, log_dir, name, interval=1, dummies=True):
+        self.enabled = enabled
+        self.interval = interval
+        self.cache = {}
+        if self.enabled:
+            self.summary_writer = SummaryWriter(
+                log_dir=os.path.join(log_dir, name),
+                flush_secs=120, max_queue=200)
+            atexit.register(self.summary_writer.close)
+            if dummies:
+                for key in ('aaa', 'zzz'):
+                    self.summary_writer.add_scalar(key, 0.0, 1)
+
+    def log(self, step, data):
+        for k, v in data.items():
+            self.log_value(step, k, v.item() if type(v) is torch.Tensor else v)
+
+    def log_value(self, step, key, val, stat='mean'):
+        if self.enabled:
+            if key not in self.cache:
+                self.cache[key] = []
+            self.cache[key].append(val)
+            if len(self.cache[key]) == self.interval:
+                agg_val = getattr(np, stat)(self.cache[key])
+                self.summary_writer.add_scalar(key, agg_val, step)
+                del self.cache[key]
+
+    def log_grads(self, step, model):
+        if self.enabled:
+            norms = [p.grad.norm().item() for p in model.parameters()
+                     if p.grad is not None]
+            for stat in ('max', 'min', 'mean'):
+                self.log_value(step, f'grad_{stat}', getattr(np, stat)(norms),
+                               stat=stat)
+
+
+def unique_log_fpath(log_fpath):
+
+    if not os.path.isfile(log_fpath):
+        return log_fpath
+
+    # Avoid overwriting old logs
+    saved = sorted([int(re.search('\.(\d+)', f).group(1))
+                    for f in glob.glob(f'{log_fpath}.*')])
+
+    log_num = (saved[-1] if saved else 0) + 1
+    return f'{log_fpath}.{log_num}'
+
+
+def stdout_step_format(step):
+    if isinstance(step, str):
+        return step
+    fields = []
+    if len(step) > 0:
+        fields.append("epoch {:>4}".format(step[0]))
+    if len(step) > 1:
+        fields.append("iter {:>4}".format(step[1]))
+    if len(step) > 2:
+        fields[-1] += "/{}".format(step[2])
+    return " | ".join(fields)
+
+
+def stdout_metric_format(metric, metadata, value):
+    name = metadata.get("name", metric + " : ")
+    unit = metadata.get("unit", None)
+    format = f'{{{metadata.get("format", "")}}}'
+    fields = [name, format.format(value) if value is not None else value, unit]
+    fields = [f for f in fields if f is not None]
+    return "| " + " ".join(fields)
+
+
+def init_log(args):
+    enabled = (args.local_rank == 0)
+    if enabled:
+        fpath = args.log_file or os.path.join(args.output_dir, 'nvlog.json')
+        backends = [JSONStreamBackend(Verbosity.DEFAULT,
+                                      unique_log_fpath(fpath)),
+                    StdOutBackend(Verbosity.VERBOSE,
+                                  step_format=stdout_step_format,
+                                  metric_format=stdout_metric_format)]
+    else:
+        backends = []
+
+    dllogger.init(backends=backends)
+    dllogger.metadata("train_lrate", {"name": "lrate", "format": ":>3.2e"})
+
+    for id_, pref in [('train', ''), ('train_avg', 'avg train '),
+                      ('dev', '  avg dev '), ('dev_ema', '  EMA dev ')]:
+
+        dllogger.metadata(f"{id_}_loss",
+                          {"name": f"{pref}loss", "format": ":>7.2f"})
+
+        dllogger.metadata(f"{id_}_wer",
+                          {"name": f"{pref}wer", "format": ":>6.2f"})
+
+        dllogger.metadata(f"{id_}_throughput",
+                          {"name": f"{pref}utts/s", "format": ":>5.0f"})
+
+        dllogger.metadata(f"{id_}_took",
+                          {"name": "took", "unit": "s", "format": ":>5.2f"})
+
+    tb_subsets = ['train', 'dev', 'dev_ema'] if args.ema else ['train', 'dev']
+    global tb_loggers
+    tb_loggers = {s: TBLogger(enabled, args.output_dir, name=s)
+                  for s in tb_subsets}
+
+    log_parameters(vars(args), tb_subset='train')
+
+
+def log(step, tb_total_steps=None, subset='train', data={}):
+
+    if tb_total_steps is not None:
+        tb_loggers[subset].log(tb_total_steps, data)
+
+    if subset != '':
+        data = {f'{subset}_{key}': v for key,v in data.items()}
+    dllogger.log(step, data=data)
+
+
+def log_grads_tb(tb_total_steps, grads, tb_subset='train'):
+    tb_loggers[tb_subset].log_grads(tb_total_steps, grads)
+
+
+def log_parameters(data, verbosity=0, tb_subset=None):
+    for k,v in data.items():
+        dllogger.log(step="PARAMETER", data={k:v}, verbosity=verbosity)
+
+    if tb_subset is not None and tb_loggers[tb_subset].enabled:
+        tb_data = {k:v for k,v in data.items()
+                   if type(v) in (str, bool, int, float)}
+        tb_loggers[tb_subset].summary_writer.add_hparams(tb_data, {})
+
+
+def flush_log():
+    dllogger.flush()
+    for tbl in tb_loggers.values():
+        if tbl.enabled:
+            tbl.summary_writer.flush()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/LICENSE b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/LICENSE
new file mode 100644
index 0000000000..4ad4ed1d5e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/LICENSE
@@ -0,0 +1,19 @@
+Copyright (c) 2017 Keith Ito
+
+Permission is hereby granted, free of charge, to any person obtaining a copy
+of this software and associated documentation files (the "Software"), to deal
+in the Software without restriction, including without limitation the rights
+to use, copy, modify, merge, publish, distribute, sublicense, and/or sell
+copies of the Software, and to permit persons to whom the Software is
+furnished to do so, subject to the following conditions:
+
+The above copyright notice and this permission notice shall be included in
+all copies or substantial portions of the Software.
+
+THE SOFTWARE IS PROVIDED "AS IS", WITHOUT WARRANTY OF ANY KIND, EXPRESS OR
+IMPLIED, INCLUDING BUT NOT LIMITED TO THE WARRANTIES OF MERCHANTABILITY,
+FITNESS FOR A PARTICULAR PURPOSE AND NONINFRINGEMENT. IN NO EVENT SHALL THE
+AUTHORS OR COPYRIGHT HOLDERS BE LIABLE FOR ANY CLAIM, DAMAGES OR OTHER
+LIABILITY, WHETHER IN AN ACTION OF CONTRACT, TORT OR OTHERWISE, ARISING FROM,
+OUT OF OR IN CONNECTION WITH THE SOFTWARE OR THE USE OR OTHER DEALINGS IN
+THE SOFTWARE.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/__init__.py
new file mode 100644
index 0000000000..4901823853
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/__init__.py
@@ -0,0 +1,32 @@
+# Copyright (c) 2017 Keith Ito
+""" from https://github.com/keithito/tacotron """
+import re
+import string
+from . import cleaners
+
+def _clean_text(text, cleaner_names, *args):
+    for name in cleaner_names:
+        cleaner = getattr(cleaners, name)
+        if not cleaner:
+            raise Exception('Unknown cleaner: %s' % name)
+        text = cleaner(text, *args)
+    return text
+
+
+def punctuation_map(labels):
+    # Punctuation to remove
+    punctuation = string.punctuation
+    punctuation = punctuation.replace("+", "")
+    punctuation = punctuation.replace("&", "")
+    # TODO We might also want to consider:
+    # @ -> at
+    # # -> number, pound, hashtag
+    # ~ -> tilde
+    # _ -> underscore
+    # % -> percent
+    # If a punctuation symbol is inside our vocab, we do not remove from text
+    for l in labels:
+        punctuation = punctuation.replace(l, "")
+    # Turn all punctuation to whitespace
+    table = str.maketrans(punctuation, " " * len(punctuation))
+    return table
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/cleaners.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/cleaners.py
new file mode 100644
index 0000000000..a99db1a625
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/cleaners.py
@@ -0,0 +1,107 @@
+# Copyright (c) 2017 Keith Ito
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+""" from https://github.com/keithito/tacotron 
+Modified to add puncturation removal
+"""
+
+'''
+Cleaners are transformations that run over the input text at both training and eval time.
+
+Cleaners can be selected by passing a comma-delimited list of cleaner names as the "cleaners"
+hyperparameter. Some cleaners are English-specific. You'll typically want to use:
+    1. "english_cleaners" for English text
+    2. "transliteration_cleaners" for non-English text that can be transliterated to ASCII using
+         the Unidecode library (https://pypi.python.org/pypi/Unidecode)
+    3. "basic_cleaners" if you do not want to transliterate (in this case, you should also update
+         the symbols in symbols.py to match your data).
+
+'''
+
+import re
+from unidecode import unidecode
+from .numbers import normalize_numbers
+
+# Regular expression matching whitespace:
+_whitespace_re = re.compile(r'\s+')
+
+# List of (regular expression, replacement) pairs for abbreviations:
+_abbreviations = [(re.compile('\\b%s\\.' % x[0], re.IGNORECASE), x[1]) for x in [
+    ('mrs', 'misess'),
+    ('mr', 'mister'),
+    ('dr', 'doctor'),
+    ('st', 'saint'),
+    ('co', 'company'),
+    ('jr', 'junior'),
+    ('maj', 'major'),
+    ('gen', 'general'),
+    ('drs', 'doctors'),
+    ('rev', 'reverend'),
+    ('lt', 'lieutenant'),
+    ('hon', 'honorable'),
+    ('sgt', 'sergeant'),
+    ('capt', 'captain'),
+    ('esq', 'esquire'),
+    ('ltd', 'limited'),
+    ('col', 'colonel'),
+    ('ft', 'fort'),
+]]
+
+def expand_abbreviations(text):
+    for regex, replacement in _abbreviations:
+        text = re.sub(regex, replacement, text)
+    return text
+
+def expand_numbers(text):
+    return normalize_numbers(text)
+
+def lowercase(text):
+    return text.lower()
+
+def collapse_whitespace(text):
+    return re.sub(_whitespace_re, ' ', text)
+
+def convert_to_ascii(text):
+    return unidecode(text)
+
+def remove_punctuation(text, table):
+    text = text.translate(table)
+    text = re.sub(r'&', " and ", text)
+    text = re.sub(r'\+', " plus ", text)
+    return text
+
+def basic_cleaners(text):
+    '''Basic pipeline that lowercases and collapses whitespace without transliteration.'''
+    text = lowercase(text)
+    text = collapse_whitespace(text)
+    return text
+
+def transliteration_cleaners(text):
+    '''Pipeline for non-English text that transliterates to ASCII.'''
+    text = convert_to_ascii(text)
+    text = lowercase(text)
+    text = collapse_whitespace(text)
+    return text
+
+def english_cleaners(text, table=None):
+    '''Pipeline for English text, including number and abbreviation expansion.'''
+    text = convert_to_ascii(text)
+    text = lowercase(text)
+    text = expand_numbers(text)
+    text = expand_abbreviations(text)
+    if table is not None:
+        text = remove_punctuation(text, table)
+    text = collapse_whitespace(text)
+    return text
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/numbers.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/numbers.py
new file mode 100644
index 0000000000..46ce110676
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/numbers.py
@@ -0,0 +1,99 @@
+# Copyright (c) 2017 Keith Ito
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+""" from https://github.com/keithito/tacotron 
+Modifed to add support for time and slight tweaks to _expand_number
+"""
+
+import inflect
+import re
+
+
+_inflect = inflect.engine()
+_comma_number_re = re.compile(r'([0-9][0-9\,]+[0-9])')
+_decimal_number_re = re.compile(r'([0-9]+\.[0-9]+)')
+_pounds_re = re.compile(r'£([0-9\,]*[0-9]+)')
+_dollars_re = re.compile(r'\$([0-9\.\,]*[0-9]+)')
+_ordinal_re = re.compile(r'[0-9]+(st|nd|rd|th)')
+_number_re = re.compile(r'[0-9]+')
+_time_re = re.compile(r'([0-9]{1,2}):([0-9]{2})')
+
+
+def _remove_commas(m):
+    return m.group(1).replace(',', '')
+
+
+def _expand_decimal_point(m):
+    return m.group(1).replace('.', ' point ')
+
+
+def _expand_dollars(m):
+    match = m.group(1)
+    parts = match.split('.')
+    if len(parts) > 2:
+        return match + ' dollars'  # Unexpected format
+    dollars = int(parts[0]) if parts[0] else 0
+    cents = int(parts[1]) if len(parts) > 1 and parts[1] else 0
+    if dollars and cents:
+        dollar_unit = 'dollar' if dollars == 1 else 'dollars'
+        cent_unit = 'cent' if cents == 1 else 'cents'
+        return '%s %s, %s %s' % (dollars, dollar_unit, cents, cent_unit)
+    elif dollars:
+        dollar_unit = 'dollar' if dollars == 1 else 'dollars'
+        return '%s %s' % (dollars, dollar_unit)
+    elif cents:
+        cent_unit = 'cent' if cents == 1 else 'cents'
+        return '%s %s' % (cents, cent_unit)
+    else:
+        return 'zero dollars'
+
+
+def _expand_ordinal(m):
+    return _inflect.number_to_words(m.group(0))
+
+
+def _expand_number(m):
+    if int(m.group(0)[0]) == 0:
+        return _inflect.number_to_words(m.group(0), andword='', group=1)
+    num = int(m.group(0))
+    if num > 1000 and num < 3000:
+        if num == 2000:
+            return 'two thousand'
+        elif num > 2000 and num < 2010:
+            return 'two thousand ' + _inflect.number_to_words(num % 100)
+        elif num % 100 == 0:
+            return _inflect.number_to_words(num // 100) + ' hundred'
+        else:
+            return _inflect.number_to_words(num, andword='', zero='oh', group=2).replace(', ', ' ')
+    # Add check for number phones and other large numbers
+    elif num > 1000000000 and num % 10000 != 0:
+        return _inflect.number_to_words(num, andword='', group=1)
+    else:
+        return _inflect.number_to_words(num, andword='')
+
+def _expand_time(m):
+    mins = int(m.group(2))
+    if mins == 0:
+        return _inflect.number_to_words(m.group(1))
+    return " ".join([_inflect.number_to_words(m.group(1)), _inflect.number_to_words(m.group(2))])
+
+def normalize_numbers(text):
+    text = re.sub(_comma_number_re, _remove_commas, text)
+    text = re.sub(_pounds_re, r'\1 pounds', text)
+    text = re.sub(_dollars_re, _expand_dollars, text)
+    text = re.sub(_decimal_number_re, _expand_decimal_point, text)
+    text = re.sub(_ordinal_re, _expand_ordinal, text)
+    text = re.sub(_number_re, _expand_number, text)
+    text = re.sub(_time_re, _expand_time, text)
+    return text
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/symbols.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/symbols.py
new file mode 100644
index 0000000000..24efedf8da
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/text/symbols.py
@@ -0,0 +1,19 @@
+# Copyright (c) 2017 Keith Ito
+""" from https://github.com/keithito/tacotron """
+
+'''
+Defines the set of symbols used in text input to the model.
+
+The default is a set of ASCII characters that works well for English or text that has been run through Unidecode. For other data, you can modify _characters. See TRAINING_DATA.md for details. '''
+from . import cmudict
+
+_pad        = '_'
+_punctuation = '!\'(),.:;? '
+_special = '-'
+_letters = 'ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz'
+
+# Prepend "@" to ARPAbet symbols to ensure uniqueness (some are the same as uppercase letters):
+_arpabet = ['@' + s for s in cmudict.valid_symbols]
+
+# Export all symbols:
+symbols = [_pad] + list(_special) + list(_punctuation) + list(_letters) + _arpabet
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/train.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/train.py
new file mode 100644
index 0000000000..2fec7f56ed
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/train.py
@@ -0,0 +1,516 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import argparse
+import copy
+import os
+import random
+import time
+
+try:
+    import nvidia_dlprof_pytorch_nvtx as pyprof
+except ModuleNotFoundError:
+    import pyprof
+
+import torch
+import numpy as np
+import torch.cuda.profiler as profiler
+import torch.distributed as dist
+from apex import amp
+from apex.parallel import DistributedDataParallel
+
+from common import helpers
+# from common.dali.data_loader import DaliDataLoader
+from common.dataset import AudioDataset, get_data_loader
+from common.features import BaseFeatures, FilterbankFeatures
+from common.helpers import (Checkpointer, greedy_wer, num_weights, print_once,
+                            process_evaluation_epoch)
+from common.optimizers import AdamW, lr_policy, Novograd
+from common.tb_dllogger import flush_log, init_log, log
+from common.utils import BenchmarkStats
+from jasper import config
+from jasper.model import CTCLossNM, GreedyCTCDecoder, Jasper
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Jasper')
+
+    training = parser.add_argument_group('training setup')
+    training.add_argument('--epochs', default=400, type=int,
+                          help='Number of epochs for the entire training; influences the lr schedule')
+    training.add_argument("--warmup_epochs", default=0, type=int,
+                          help='Initial epochs of increasing learning rate')
+    training.add_argument("--hold_epochs", default=0, type=int,
+                          help='Constant max learning rate epochs after warmup')
+    training.add_argument('--epochs_this_job', default=0, type=int,
+                          help=('Run for a number of epochs with no effect on the lr schedule.'
+                                'Useful for re-starting the training.'))
+    training.add_argument('--cudnn_benchmark', action='store_true', default=True,
+                          help='Enable cudnn benchmark')
+    training.add_argument('--amp', '--fp16', action='store_true', default=False,
+                          help='Use mixed precision training')
+    training.add_argument('--seed', default=42, type=int, help='Random seed')
+    training.add_argument('--local_rank', default=os.getenv('LOCAL_RANK', 0),
+                          type=int, help='GPU id used for distributed training')
+    training.add_argument('--pre_allocate_range', default=None, type=int, nargs=2,
+                          help='Warmup with batches of length [min, max] before training')
+    training.add_argument('--pyprof', action='store_true', help='Enable pyprof profiling')
+
+    optim = parser.add_argument_group('optimization setup')
+    optim.add_argument('--batch_size', default=32, type=int,
+                       help='Global batch size')
+    optim.add_argument('--lr', default=1e-3, type=float,
+                       help='Peak learning rate')
+    optim.add_argument("--min_lr", default=1e-5, type=float,
+                       help='minimum learning rate')
+    optim.add_argument("--lr_policy", default='exponential', type=str,
+                       choices=['exponential', 'legacy'], help='lr scheduler')
+    optim.add_argument("--lr_exp_gamma", default=0.99, type=float,
+                       help='gamma factor for exponential lr scheduler')
+    optim.add_argument('--weight_decay', default=1e-3, type=float,
+                       help='Weight decay for the optimizer')
+    optim.add_argument('--grad_accumulation_steps', default=1, type=int,
+                       help='Number of accumulation steps')
+    optim.add_argument('--optimizer', default='novograd', type=str,
+                       choices=['novograd', 'adamw'], help='Optimization algorithm')
+    optim.add_argument('--ema', type=float, default=0.0,
+                       help='Discount factor for exp averaging of model weights')
+
+    io = parser.add_argument_group('feature and checkpointing setup')
+    io.add_argument('--dali_device', type=str, choices=['none', 'cpu', 'gpu'],
+                    default='none', help='Use DALI pipeline for fast data processing')
+    io.add_argument('--resume', action='store_true',
+                    help='Try to resume from last saved checkpoint.')
+    io.add_argument('--ckpt', default=None, type=str,
+                    help='Path to a checkpoint for resuming training')
+    io.add_argument('--save_frequency', default=10, type=int,
+                    help='Checkpoint saving frequency in epochs')
+    io.add_argument('--keep_milestones', default=[100, 200, 300], type=int, nargs='+',
+                    help='Milestone checkpoints to keep from removing')
+    io.add_argument('--save_best_from', default=380, type=int,
+                    help='Epoch on which to begin tracking best checkpoint (dev WER)')
+    io.add_argument('--eval_frequency', default=200, type=int,
+                    help='Number of steps between evaluations on dev set')
+    io.add_argument('--log_frequency', default=25, type=int,
+                    help='Number of steps between printing training stats')
+    io.add_argument('--prediction_frequency', default=100, type=int,
+                    help='Number of steps between printing sample decodings')
+    io.add_argument('--model_config', type=str, required=True,
+                    help='Path of the model configuration file')
+    io.add_argument('--train_manifests', type=str, required=True, nargs='+',
+                    help='Paths of the training dataset manifest file')
+    io.add_argument('--val_manifests', type=str, required=True, nargs='+',
+                    help='Paths of the evaluation datasets manifest files')
+    io.add_argument('--dataset_dir', required=True, type=str,
+                    help='Root dir of dataset')
+    io.add_argument('--output_dir', type=str, required=True,
+                    help='Directory for logs and checkpoints')
+    io.add_argument('--log_file', type=str, default=None,
+                    help='Path to save the training logfile.')
+    io.add_argument('--benchmark_epochs_num', type=int, default=1,
+                    help='Number of epochs accounted in final average throughput.')
+    io.add_argument('--override_config', type=str, action='append',
+                    help='Overrides a value from a config .yaml.'
+                         ' Syntax: `--override_config nested.config.key=val`.')
+    return parser.parse_args()
+
+
+def reduce_tensor(tensor, num_gpus):
+    rt = tensor.clone()
+    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+    return rt.true_divide(num_gpus)
+
+
+def apply_ema(model, ema_model, decay):
+    if not decay:
+        return
+
+    sd = getattr(model, 'module', model).state_dict()
+    for k, v in ema_model.state_dict().items():
+        v.copy_(decay * v + (1 - decay) * sd[k])
+
+
+@torch.no_grad()
+def evaluate(epoch, step, val_loader, val_feat_proc, labels, model,
+             ema_model, ctc_loss, greedy_decoder, use_amp, use_dali=False):
+
+    for model, subset in [(model, 'dev'), (ema_model, 'dev_ema')]:
+        if model is None:
+            continue
+
+        model.eval()
+        start_time = time.time()
+        agg = {'losses': [], 'preds': [], 'txts': []}
+
+        for batch in val_loader:
+            if use_dali:
+                # with DALI, the data is already on GPU
+                feat, feat_lens, txt, txt_lens = batch
+                if val_feat_proc is not None:
+                    feat, feat_lens = val_feat_proc(feat, feat_lens, use_amp)
+            else:
+                batch = [t.npu(non_blocking=True) for t in batch]
+                audio, audio_lens, txt, txt_lens = batch
+                feat, feat_lens = val_feat_proc(audio, audio_lens, use_amp)
+
+            log_probs, enc_lens = model.forward(feat, feat_lens)
+            loss = ctc_loss(log_probs, txt, enc_lens, txt_lens)
+            pred = greedy_decoder(log_probs)
+
+            agg['losses'] += helpers.gather_losses([loss])
+            agg['preds'] += helpers.gather_predictions([pred], labels)
+            agg['txts'] += helpers.gather_transcripts([txt], [txt_lens], labels)
+
+        wer, loss = process_evaluation_epoch(agg)
+        log((epoch,), step, subset, {'loss': loss, 'wer': 100.0 * wer,
+                                     'took': time.time() - start_time})
+        model.train()
+    return wer
+
+
+def main():
+    args = parse_args()
+
+    assert(torch.npu.is_available())
+    assert args.prediction_frequency % args.log_frequency == 0
+
+    torch.backends.cudnn.benchmark = args.cudnn_benchmark
+
+    # set up distributed training
+    multi_gpu = False
+    if multi_gpu:
+        torch.cuda.set_device(args.local_rank)
+        dist.init_process_group(backend='nccl', init_method='env://')
+        world_size = dist.get_world_size()
+        print_once(f'Distributed training with {world_size} GPUs\n')
+    else:
+        world_size = 1
+
+    torch.manual_seed(args.seed + args.local_rank)
+    np.random.seed(args.seed + args.local_rank)
+    random.seed(args.seed + args.local_rank)
+
+    init_log(args)
+
+    cfg = config.load(args.model_config)
+    config.apply_config_overrides(cfg, args)
+
+    symbols = helpers.add_ctc_blank(cfg['labels'])
+
+    assert args.grad_accumulation_steps >= 1
+    assert args.batch_size % args.grad_accumulation_steps == 0
+    batch_size = args.batch_size // args.grad_accumulation_steps
+
+    print_once('Setting up datasets...')
+    train_dataset_kw, train_features_kw = config.input(cfg, 'train')
+    val_dataset_kw, val_features_kw = config.input(cfg, 'val')
+
+    # use_dali = args.dali_device in ('cpu', 'gpu')
+    use_dali = False
+    if use_dali:
+        assert train_dataset_kw['ignore_offline_speed_perturbation'], \
+            "DALI doesn't support offline speed perturbation"
+
+        # pad_to_max_duration is not supported by DALI - have simple padders
+        if train_features_kw['pad_to_max_duration']:
+            train_feat_proc = BaseFeatures(
+                pad_align=train_features_kw['pad_align'],
+                pad_to_max_duration=True,
+                max_duration=train_features_kw['max_duration'],
+                sample_rate=train_features_kw['sample_rate'],
+                window_size=train_features_kw['window_size'],
+                window_stride=train_features_kw['window_stride'])
+            train_features_kw['pad_to_max_duration'] = False
+        else:
+            train_feat_proc = None
+
+        if val_features_kw['pad_to_max_duration']:
+            val_feat_proc = BaseFeatures(
+                pad_align=val_features_kw['pad_align'],
+                pad_to_max_duration=True,
+                max_duration=val_features_kw['max_duration'],
+                sample_rate=val_features_kw['sample_rate'],
+                window_size=val_features_kw['window_size'],
+                window_stride=val_features_kw['window_stride'])
+            val_features_kw['pad_to_max_duration'] = False
+        else:
+            val_feat_proc = None
+
+        train_loader = DaliDataLoader(gpu_id=args.local_rank,
+                                      dataset_path=args.dataset_dir,
+                                      config_data=train_dataset_kw,
+                                      config_features=train_features_kw,
+                                      json_names=args.train_manifests,
+                                      batch_size=batch_size,
+                                      grad_accumulation_steps=args.grad_accumulation_steps,
+                                      pipeline_type="train",
+                                      device_type=args.dali_device,
+                                      symbols=symbols)
+
+        val_loader = DaliDataLoader(gpu_id=args.local_rank,
+                                    dataset_path=args.dataset_dir,
+                                    config_data=val_dataset_kw,
+                                    config_features=val_features_kw,
+                                    json_names=args.val_manifests,
+                                    batch_size=batch_size,
+                                    pipeline_type="val",
+                                    device_type=args.dali_device,
+                                    symbols=symbols)
+    else:
+        train_dataset_kw, train_features_kw = config.input(cfg, 'train')
+        train_dataset = AudioDataset(args.dataset_dir,
+                                     args.train_manifests,
+                                     symbols,
+                                     **train_dataset_kw)
+        train_loader = get_data_loader(train_dataset,
+                                       batch_size,
+                                       multi_gpu=multi_gpu,
+                                       shuffle=True,
+                                       num_workers=4)
+        train_feat_proc = FilterbankFeatures(**train_features_kw)
+
+        val_dataset_kw, val_features_kw = config.input(cfg, 'val')
+        val_dataset = AudioDataset(args.dataset_dir,
+                                   args.val_manifests,
+                                   symbols,
+                                   **val_dataset_kw)
+        val_loader = get_data_loader(val_dataset,
+                                     batch_size,
+                                     multi_gpu=multi_gpu,
+                                     shuffle=False,
+                                     num_workers=4,
+                                     drop_last=False)
+        val_feat_proc = FilterbankFeatures(**val_features_kw)
+
+        dur = train_dataset.duration / 3600
+        dur_f = train_dataset.duration_filtered / 3600
+        nsampl = len(train_dataset)
+        print_once(f'Training samples: {nsampl} ({dur:.1f}h, '
+                   f'filtered {dur_f:.1f}h)')
+
+    if train_feat_proc is not None:
+        train_feat_proc.cpu()
+    if val_feat_proc is not None:
+        val_feat_proc.cpu()
+
+    steps_per_epoch = len(train_loader) // args.grad_accumulation_steps
+
+    # set up the model
+    model = Jasper(encoder_kw=config.encoder(cfg),
+                   decoder_kw=config.decoder(cfg, n_classes=len(symbols)))
+    model.cpu()
+    ctc_loss = CTCLossNM(n_classes=len(symbols))
+    greedy_decoder = GreedyCTCDecoder()
+
+    print_once(f'Model size: {num_weights(model) / 10**6:.1f}M params\n')
+
+    # optimization
+    kw = {'lr': args.lr, 'weight_decay': args.weight_decay}
+    if args.optimizer == "novograd":
+        optimizer = Novograd(model.parameters(), **kw)
+    elif args.optimizer == "adamw":
+        optimizer = AdamW(model.parameters(), **kw)
+    else:
+        raise ValueError(f'Invalid optimizer "{args.optimizer}"')
+
+    adjust_lr = lambda step, epoch, optimizer: lr_policy(
+        step, epoch, args.lr, optimizer, steps_per_epoch=steps_per_epoch,
+        warmup_epochs=args.warmup_epochs, hold_epochs=args.hold_epochs,
+        num_epochs=args.epochs, policy=args.lr_policy, min_lr=args.min_lr,
+        exp_gamma=args.lr_exp_gamma)
+
+    if args.amp:
+        model, optimizer = amp.initialize(
+            min_loss_scale=1.0, models=model, optimizers=optimizer,
+            opt_level='O1', max_loss_scale=512.0)
+
+    if args.ema > 0:
+        ema_model = copy.deepcopy(model)
+    else:
+        ema_model = None
+
+    if multi_gpu:
+        model = DistributedDataParallel(model)
+
+    if args.pyprof:
+        pyprof.init(enable_function_stack=True)
+
+    # load checkpoint
+    meta = {'best_wer': 10**6, 'start_epoch': 0}
+    checkpointer = Checkpointer(args.output_dir, 'Jasper',
+                                args.keep_milestones, args.amp)
+    if args.resume:
+        args.ckpt = checkpointer.last_checkpoint() or args.ckpt
+
+    if args.ckpt is not None:
+        checkpointer.load(args.ckpt, model, ema_model, optimizer, meta)
+
+    start_epoch = meta['start_epoch']
+    best_wer = meta['best_wer']
+    epoch = 1
+    step = start_epoch * steps_per_epoch + 1
+
+    if args.pyprof:
+        torch.autograd.profiler.emit_nvtx().__enter__()
+        profiler.start()
+
+    # training loop
+    model.train()
+
+    # pre-allocate
+    if args.pre_allocate_range is not None:
+        n_feats = train_features_kw['n_filt']
+        pad_align = train_features_kw['pad_align']
+        a, b = args.pre_allocate_range
+        for n_frames in range(a, b + pad_align, pad_align):
+            print_once(f'Pre-allocation ({batch_size}x{n_feats}x{n_frames})...')
+
+            feat = torch.randn(batch_size, n_feats, n_frames, device='cpu')
+            feat_lens = torch.ones(batch_size, device='cpu').fill_(n_frames)
+            txt = torch.randint(high=len(symbols)-1, size=(batch_size, 100),
+                                device='cpu')
+            txt_lens = torch.ones(batch_size, device='cpu').fill_(100)
+            log_probs, enc_lens = model(feat, feat_lens)
+            del feat
+            loss = ctc_loss(log_probs, txt, enc_lens, txt_lens)
+            loss.backward()
+            model.zero_grad()
+
+    bmark_stats = BenchmarkStats()
+
+    for epoch in range(start_epoch + 1, args.epochs + 1):
+        if multi_gpu and not use_dali:
+            train_loader.sampler.set_epoch(epoch)
+
+        epoch_utts = 0
+        epoch_loss = 0
+        accumulated_batches = 0
+        epoch_start_time = time.time()
+
+        for batch in train_loader:
+
+            if accumulated_batches == 0:
+                adjust_lr(step, epoch, optimizer)
+                optimizer.zero_grad()
+                step_loss = 0
+                step_utts = 0
+                step_start_time = time.time()
+
+            if use_dali:
+                # with DALI, the data is already on GPU
+                feat, feat_lens, txt, txt_lens = batch
+                if train_feat_proc is not None:
+                    feat, feat_lens = train_feat_proc(feat, feat_lens, args.amp)
+            else:
+                batch = [t.npu(non_blocking=True) for t in batch]
+                audio, audio_lens, txt, txt_lens = batch
+                feat, feat_lens = train_feat_proc(audio, audio_lens, args.amp)
+            print("feat",feat)
+            print("feat_len",feat_lens)
+            log_probs, enc_lens = model(feat, feat_lens)
+
+            loss = ctc_loss(log_probs, txt, enc_lens, txt_lens)
+            loss /= args.grad_accumulation_steps
+
+            if torch.isnan(loss).any():
+                print_once(f'WARNING: loss is NaN; skipping update')
+            else:
+                if multi_gpu:
+                    step_loss += reduce_tensor(loss.data, world_size).item()
+                else:
+                    step_loss += loss.item()
+
+                if args.amp:
+                    with amp.scale_loss(loss, optimizer) as scaled_loss:
+                        scaled_loss.backward()
+                else:
+                    loss.backward()
+                step_utts += batch[0].size(0) * world_size
+                epoch_utts += batch[0].size(0) * world_size
+                accumulated_batches += 1
+
+            if accumulated_batches % args.grad_accumulation_steps == 0:
+                epoch_loss += step_loss
+                optimizer.step()
+                apply_ema(model, ema_model, args.ema)
+
+                if step % args.log_frequency == 0:
+                    preds = greedy_decoder(log_probs)
+                    wer, pred_utt, ref = greedy_wer(preds, txt, txt_lens, symbols)
+
+                    if step % args.prediction_frequency == 0:
+                        print_once(f'  Decoded:   {pred_utt[:90]}')
+                        print_once(f'  Reference: {ref[:90]}')
+
+                    step_time = time.time() - step_start_time
+                    log((epoch, step % steps_per_epoch or steps_per_epoch, steps_per_epoch),
+                        step, 'train',
+                        {'loss': step_loss,
+                         'wer': 100.0 * wer,
+                         'throughput': step_utts / step_time,
+                         'took': step_time,
+                         'lrate': optimizer.param_groups[0]['lr']})
+
+                step_start_time = time.time()
+
+                if step % args.eval_frequency == 0:
+                    wer = evaluate(epoch, step, val_loader, val_feat_proc,
+                                   symbols, model, ema_model, ctc_loss,
+                                   greedy_decoder, args.amp, use_dali)
+
+                    if wer < best_wer and epoch >= args.save_best_from:
+                        checkpointer.save(model, ema_model, optimizer, epoch,
+                                          step, best_wer, is_best=True)
+                        best_wer = wer
+
+                step += 1
+                accumulated_batches = 0
+                # end of step
+
+            # DALI iterator need to be exhausted;
+            # if not using DALI, simulate drop_last=True with grad accumulation
+            if not use_dali and step > steps_per_epoch * epoch:
+                break
+
+        epoch_time = time.time() - epoch_start_time
+        epoch_loss /= steps_per_epoch
+        log((epoch,), None, 'train_avg', {'throughput': epoch_utts / epoch_time,
+                                          'took': epoch_time,
+                                          'loss': epoch_loss})
+        bmark_stats.update(epoch_utts, epoch_time, epoch_loss)
+
+        if epoch % args.save_frequency == 0 or epoch in args.keep_milestones:
+            checkpointer.save(model, ema_model, optimizer, epoch, step, best_wer)
+
+        if 0 < args.epochs_this_job <= epoch - start_epoch:
+            print_once(f'Finished after {args.epochs_this_job} epochs.')
+            break
+        # end of epoch
+
+    if args.pyprof:
+        profiler.stop()
+        torch.autograd.profiler.emit_nvtx().__exit__(None, None, None)
+
+    log((), None, 'train_avg', bmark_stats.get(args.benchmark_epochs_num))
+
+    if epoch == args.epochs:
+        evaluate(epoch, step, val_loader, val_feat_proc, symbols, model,
+                 ema_model, ctc_loss, greedy_decoder, args.amp, use_dali)
+
+        checkpointer.save(model, ema_model, optimizer, epoch, step, best_wer)
+    flush_log()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/utils.py
new file mode 100644
index 0000000000..2bd7986f19
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/common/utils.py
@@ -0,0 +1,20 @@
+import numpy as np
+
+
+class BenchmarkStats:
+    """ Tracks statistics used for benchmarking. """
+    def __init__(self):
+        self.utts = []
+        self.times = []
+        self.losses = []
+
+    def update(self, utts, times, losses):
+        self.utts.append(utts)
+        self.times.append(times)
+        self.losses.append(losses)
+
+    def get(self, n_epochs):
+        throughput = sum(self.utts[-n_epochs:]) / sum(self.times[-n_epochs:])
+
+        return {'throughput': throughput, 'benchmark_epochs_num': n_epochs,
+                'loss': np.mean(self.losses[-n_epochs:])}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speca.yaml b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speca.yaml
new file mode 100644
index 0000000000..b0c0d5b9c4
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speca.yaml
@@ -0,0 +1,139 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+name: "Jasper"
+labels: [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+         "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+input_val:
+  audio_dataset: &val_dataset
+    sample_rate: &sample_rate 16000
+    trim_silence: true
+    normalize_transcripts: true
+
+  filterbank_features: &val_features
+    normalize: per_feature
+    sample_rate: *sample_rate
+    window_size: 0.02
+    window_stride: 0.01
+    window: hann
+    n_filt: &n_filt 64
+    n_fft: 512
+    frame_splicing: &frame_splicing 1
+    dither: 0.00001
+    pad_align: 16
+
+# For training we keep samples < 16.7s and apply augmentation
+input_train:
+  audio_dataset:
+    <<: *val_dataset
+    max_duration: 16.7
+    ignore_offline_speed_perturbation: true
+
+  filterbank_features:
+    <<: *val_features
+    max_duration: 16.7
+
+    spec_augment:
+      freq_masks: 2
+      max_freq: 20
+      time_masks: 2
+      max_time: 75
+
+jasper:
+  encoder:
+    init: xavier_uniform
+    in_feats: *n_filt
+    frame_splicing: *frame_splicing
+    activation: relu
+    use_conv_masks: true
+    blocks:
+    - &Conv1
+      filters: 256
+      repeat: 1
+      kernel_size: [11]
+      stride: [2]
+      dilation: [1]
+      dropout: 0.2
+      residual: false
+    - &B1
+      filters: 256
+      repeat: 5
+      kernel_size: [11]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B1
+    - &B2
+      filters: 384
+      repeat: 5
+      kernel_size: [13]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B2
+    - &B3
+      filters: 512
+      repeat: 5
+      kernel_size: [17]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B3
+    - &B4
+      filters: 640
+      repeat: 5
+      kernel_size: [21]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B4
+    - &B5
+      filters: 768
+      repeat: 5
+      kernel_size: [25]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B5
+    - &Conv2
+      filters: 896
+      repeat: 1
+      kernel_size: [29]
+      stride: [1]
+      dilation: [2]
+      dropout: 0.4
+      residual: false
+    - &Conv3
+      filters: &enc_feats 1024
+      repeat: 1
+      kernel_size: [1]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.4
+      residual: false
+
+  decoder:
+    in_feats: *enc_feats
+    init: xavier_uniform
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline.yaml b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline.yaml
new file mode 100644
index 0000000000..89c135ea71
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline.yaml
@@ -0,0 +1,139 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+name: "Jasper"
+labels: [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+         "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+input_val:
+  audio_dataset: &val_dataset
+    sample_rate: &sample_rate 16000
+    trim_silence: true
+    normalize_transcripts: true
+
+  filterbank_features: &val_features
+    normalize: per_feature
+    sample_rate: *sample_rate
+    window_size: 0.02
+    window_stride: 0.01
+    window: hann
+    n_filt: &n_filt 64
+    n_fft: 512
+    frame_splicing: &frame_splicing 1
+    dither: 0.00001
+    pad_align: 16
+
+# For training we keep samples < 16.7s and apply augmentation
+input_train:
+  audio_dataset:
+    <<: *val_dataset
+    max_duration: 16.7
+    ignore_offline_speed_perturbation: false
+
+  filterbank_features:
+    <<: *val_features
+    max_duration: 16.7
+
+    spec_augment:
+      freq_masks: 0
+      max_freq: 20
+      time_masks: 0
+      max_time: 75
+
+jasper:
+  encoder:
+    init: xavier_uniform
+    in_feats: *n_filt
+    frame_splicing: *frame_splicing
+    activation: relu
+    use_conv_masks: true
+    blocks:
+    - &Conv1
+      filters: 256
+      repeat: 1
+      kernel_size: [11]
+      stride: [2]
+      dilation: [1]
+      dropout: 0.2
+      residual: false
+    - &B1
+      filters: 256
+      repeat: 5
+      kernel_size: [11]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B1
+    - &B2
+      filters: 384
+      repeat: 5
+      kernel_size: [13]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B2
+    - &B3
+      filters: 512
+      repeat: 5
+      kernel_size: [17]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B3
+    - &B4
+      filters: 640
+      repeat: 5
+      kernel_size: [21]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B4
+    - &B5
+      filters: 768
+      repeat: 5
+      kernel_size: [25]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B5
+    - &Conv2
+      filters: 896
+      repeat: 1
+      kernel_size: [29]
+      stride: [1]
+      dilation: [2]
+      dropout: 0.4
+      residual: false
+    - &Conv3
+      filters: &enc_feats 1024
+      repeat: 1
+      kernel_size: [1]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.4
+      residual: false
+
+  decoder:
+    in_feats: *enc_feats
+    init: xavier_uniform
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline_speca.yaml b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline_speca.yaml
new file mode 100644
index 0000000000..2c7e45818a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline_speca.yaml
@@ -0,0 +1,139 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+name: "Jasper"
+labels: [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+         "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+input_val:
+  audio_dataset: &val_dataset
+    sample_rate: &sample_rate 16000
+    trim_silence: true
+    normalize_transcripts: true
+
+  filterbank_features: &val_features
+    normalize: per_feature
+    sample_rate: *sample_rate
+    window_size: 0.02
+    window_stride: 0.01
+    window: hann
+    n_filt: &n_filt 64
+    n_fft: 512
+    frame_splicing: &frame_splicing 1
+    dither: 0.00001
+    pad_align: 16
+
+# For training we keep samples < 16.7s and apply augmentation
+input_train:
+  audio_dataset:
+    <<: *val_dataset
+    max_duration: 16.7
+    ignore_offline_speed_perturbation: false
+
+  filterbank_features:
+    <<: *val_features
+    max_duration: 16.7
+
+    spec_augment:
+      freq_masks: 2
+      max_freq: 20
+      time_masks: 2
+      max_time: 75
+
+jasper:
+  encoder:
+    init: xavier_uniform
+    in_feats: *n_filt
+    frame_splicing: *frame_splicing
+    activation: relu
+    use_conv_masks: true
+    blocks:
+    - &Conv1
+      filters: 256
+      repeat: 1
+      kernel_size: [11]
+      stride: [2]
+      dilation: [1]
+      dropout: 0.2
+      residual: false
+    - &B1
+      filters: 256
+      repeat: 5
+      kernel_size: [11]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B1
+    - &B2
+      filters: 384
+      repeat: 5
+      kernel_size: [13]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B2
+    - &B3
+      filters: 512
+      repeat: 5
+      kernel_size: [17]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B3
+    - &B4
+      filters: 640
+      repeat: 5
+      kernel_size: [21]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B4
+    - &B5
+      filters: 768
+      repeat: 5
+      kernel_size: [25]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B5
+    - &Conv2
+      filters: 896
+      repeat: 1
+      kernel_size: [29]
+      stride: [1]
+      dilation: [2]
+      dropout: 0.4
+      residual: false
+    - &Conv3
+      filters: &enc_feats 1024
+      repeat: 1
+      kernel_size: [1]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.4
+      residual: false
+
+  decoder:
+    in_feats: *enc_feats
+    init: xavier_uniform
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline_speca_nomask.yaml b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline_speca_nomask.yaml
new file mode 100644
index 0000000000..61619428a1
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-offline_speca_nomask.yaml
@@ -0,0 +1,139 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+name: "Jasper"
+labels: [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+         "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+input_val:
+  audio_dataset: &val_dataset
+    sample_rate: &sample_rate 16000
+    trim_silence: true
+    normalize_transcripts: true
+
+  filterbank_features: &val_features
+    normalize: per_feature
+    sample_rate: *sample_rate
+    window_size: 0.02
+    window_stride: 0.01
+    window: hann
+    n_filt: &n_filt 64
+    n_fft: 512
+    frame_splicing: &frame_splicing 1
+    dither: 0.00001
+    pad_align: 16
+
+# For training we keep samples < 16.7s and apply augmentation
+input_train:
+  audio_dataset:
+    <<: *val_dataset
+    max_duration: 16.7
+    ignore_offline_speed_perturbation: false
+
+  filterbank_features:
+    <<: *val_features
+    max_duration: 16.7
+
+    spec_augment:
+      freq_masks: 2
+      max_freq: 20
+      time_masks: 2
+      max_time: 75
+
+jasper:
+  encoder:
+    init: xavier_uniform
+    in_feats: *n_filt
+    frame_splicing: *frame_splicing
+    activation: relu
+    use_conv_masks: false
+    blocks:
+    - &Conv1
+      filters: 256
+      repeat: 1
+      kernel_size: [11]
+      stride: [2]
+      dilation: [1]
+      dropout: 0.2
+      residual: false
+    - &B1
+      filters: 256
+      repeat: 5
+      kernel_size: [11]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B1
+    - &B2
+      filters: 384
+      repeat: 5
+      kernel_size: [13]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B2
+    - &B3
+      filters: 512
+      repeat: 5
+      kernel_size: [17]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B3
+    - &B4
+      filters: 640
+      repeat: 5
+      kernel_size: [21]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B4
+    - &B5
+      filters: 768
+      repeat: 5
+      kernel_size: [25]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B5
+    - &Conv2
+      filters: 896
+      repeat: 1
+      kernel_size: [29]
+      stride: [1]
+      dilation: [2]
+      dropout: 0.4
+      residual: false
+    - &Conv3
+      filters: &enc_feats 1024
+      repeat: 1
+      kernel_size: [1]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.4
+      residual: false
+
+  decoder:
+    in_feats: *enc_feats
+    init: xavier_uniform
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online-discrete.yaml b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online-discrete.yaml
new file mode 100644
index 0000000000..c0c59e196b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online-discrete.yaml
@@ -0,0 +1,144 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+name: "Jasper"
+labels: [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+         "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+input_val:
+  audio_dataset: &val_dataset
+    sample_rate: &sample_rate 16000
+    trim_silence: true
+    normalize_transcripts: true
+
+  filterbank_features: &val_features
+    normalize: per_feature
+    sample_rate: *sample_rate
+    window_size: 0.02
+    window_stride: 0.01
+    window: hann
+    n_filt: &n_filt 64
+    n_fft: 512
+    frame_splicing: &frame_splicing 1
+    dither: 0.00001
+    pad_align: 16
+
+# For training we keep samples < 16.7s and apply augmentation
+input_train:
+  audio_dataset:
+    <<: *val_dataset
+    max_duration: 16.7
+    ignore_offline_speed_perturbation: true
+
+    speed_perturbation:
+      discrete: true
+      min_rate: 0.9
+      max_rate: 1.1
+
+  filterbank_features:
+    <<: *val_features
+    max_duration: 16.7
+
+    spec_augment:
+      freq_masks: 0
+      max_freq: 20
+      time_masks: 0
+      max_time: 75
+
+jasper:
+  encoder:
+    init: xavier_uniform
+    in_feats: *n_filt
+    frame_splicing: *frame_splicing
+    activation: relu
+    use_conv_masks: true
+    blocks:
+    - &Conv1
+      filters: 256
+      repeat: 1
+      kernel_size: [11]
+      stride: [2]
+      dilation: [1]
+      dropout: 0.2
+      residual: false
+    - &B1
+      filters: 256
+      repeat: 5
+      kernel_size: [11]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B1
+    - &B2
+      filters: 384
+      repeat: 5
+      kernel_size: [13]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B2
+    - &B3
+      filters: 512
+      repeat: 5
+      kernel_size: [17]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B3
+    - &B4
+      filters: 640
+      repeat: 5
+      kernel_size: [21]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B4
+    - &B5
+      filters: 768
+      repeat: 5
+      kernel_size: [25]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B5
+    - &Conv2
+      filters: 896
+      repeat: 1
+      kernel_size: [29]
+      stride: [1]
+      dilation: [2]
+      dropout: 0.4
+      residual: false
+    - &Conv3
+      filters: &enc_feats 1024
+      repeat: 1
+      kernel_size: [1]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.4
+      residual: false
+
+  decoder:
+    in_feats: *enc_feats
+    init: xavier_uniform
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online-discrete_speca.yaml b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online-discrete_speca.yaml
new file mode 100644
index 0000000000..d2491b30b7
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online-discrete_speca.yaml
@@ -0,0 +1,144 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+name: "Jasper"
+labels: [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+         "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+input_val:
+  audio_dataset: &val_dataset
+    sample_rate: &sample_rate 16000
+    trim_silence: true
+    normalize_transcripts: true
+
+  filterbank_features: &val_features
+    normalize: per_feature
+    sample_rate: *sample_rate
+    window_size: 0.02
+    window_stride: 0.01
+    window: hann
+    n_filt: &n_filt 64
+    n_fft: 512
+    frame_splicing: &frame_splicing 1
+    dither: 0.00001
+    pad_align: 16
+
+# For training we keep samples < 16.7s and apply augmentation
+input_train:
+  audio_dataset:
+    <<: *val_dataset
+    max_duration: 16.7
+    ignore_offline_speed_perturbation: true
+
+    speed_perturbation:
+      discrete: true
+      min_rate: 0.9
+      max_rate: 1.1
+
+  filterbank_features:
+    <<: *val_features
+    max_duration: 16.7
+
+    spec_augment:
+      freq_masks: 2
+      max_freq: 20
+      time_masks: 2
+      max_time: 75
+
+jasper:
+  encoder:
+    init: xavier_uniform
+    in_feats: *n_filt
+    frame_splicing: *frame_splicing
+    activation: relu
+    use_conv_masks: true
+    blocks:
+    - &Conv1
+      filters: 256
+      repeat: 1
+      kernel_size: [11]
+      stride: [2]
+      dilation: [1]
+      dropout: 0.2
+      residual: false
+    - &B1
+      filters: 256
+      repeat: 5
+      kernel_size: [11]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B1
+    - &B2
+      filters: 384
+      repeat: 5
+      kernel_size: [13]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B2
+    - &B3
+      filters: 512
+      repeat: 5
+      kernel_size: [17]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B3
+    - &B4
+      filters: 640
+      repeat: 5
+      kernel_size: [21]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B4
+    - &B5
+      filters: 768
+      repeat: 5
+      kernel_size: [25]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B5
+    - &Conv2
+      filters: 896
+      repeat: 1
+      kernel_size: [29]
+      stride: [1]
+      dilation: [2]
+      dropout: 0.4
+      residual: false
+    - &Conv3
+      filters: &enc_feats 1024
+      repeat: 1
+      kernel_size: [1]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.4
+      residual: false
+
+  decoder:
+    in_feats: *enc_feats
+    init: xavier_uniform
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online_speca.yaml b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online_speca.yaml
new file mode 100644
index 0000000000..a165af7f4a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/configs/jasper10x5dr_speedp-online_speca.yaml
@@ -0,0 +1,144 @@
+# Copyright (c) 2020, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+name: "Jasper"
+labels: [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m",
+         "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'"]
+
+input_val:
+  audio_dataset: &val_dataset
+    sample_rate: &sample_rate 16000
+    trim_silence: true
+    normalize_transcripts: true
+
+  filterbank_features: &val_features
+    normalize: per_feature
+    sample_rate: *sample_rate
+    window_size: 0.02
+    window_stride: 0.01
+    window: hann
+    n_filt: &n_filt 64
+    n_fft: 512
+    frame_splicing: &frame_splicing 1
+    dither: 0.00001
+    pad_align: 16
+
+# For training we keep samples < 16.7s and apply augmentation
+input_train:
+  audio_dataset:
+    <<: *val_dataset
+    max_duration: 16.7
+    ignore_offline_speed_perturbation: true
+
+    speed_perturbation:
+      discrete: false
+      min_rate: 0.85
+      max_rate: 1.15
+
+  filterbank_features:
+    <<: *val_features
+    max_duration: 16.7
+
+    spec_augment:
+      freq_masks: 2
+      max_freq: 20
+      time_masks: 2
+      max_time: 75
+
+jasper:
+  encoder:
+    init: xavier_uniform
+    in_feats: *n_filt
+    frame_splicing: *frame_splicing
+    activation: relu
+    use_conv_masks: true
+    blocks:
+    - &Conv1
+      filters: 256
+      repeat: 1
+      kernel_size: [11]
+      stride: [2]
+      dilation: [1]
+      dropout: 0.2
+      residual: false
+    - &B1
+      filters: 256
+      repeat: 5
+      kernel_size: [11]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B1
+    - &B2
+      filters: 384
+      repeat: 5
+      kernel_size: [13]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B2
+    - &B3
+      filters: 512
+      repeat: 5
+      kernel_size: [17]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.2
+      residual: true
+      residual_dense: true
+    - *B3
+    - &B4
+      filters: 640
+      repeat: 5
+      kernel_size: [21]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B4
+    - &B5
+      filters: 768
+      repeat: 5
+      kernel_size: [25]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.3
+      residual: true
+      residual_dense: true
+    - *B5
+    - &Conv2
+      filters: 896
+      repeat: 1
+      kernel_size: [29]
+      stride: [1]
+      dilation: [2]
+      dropout: 0.4
+      residual: false
+    - &Conv3
+      filters: &enc_feats 1024
+      repeat: 1
+      kernel_size: [1]
+      stride: [1]
+      dilation: [1]
+      dropout: 0.4
+      residual: false
+
+  decoder:
+    in_feats: *enc_feats
+    init: xavier_uniform
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/inference.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/inference.py
new file mode 100644
index 0000000000..317215e91f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/inference.py
@@ -0,0 +1,398 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import argparse
+import math
+import os
+import random
+import time
+from heapq import nlargest
+from itertools import chain, repeat
+from pathlib import Path
+from tqdm import tqdm
+
+import dllogger
+import torch
+import numpy as np
+import torch.distributed as distrib
+import torch.nn.functional as F
+from apex import amp
+from apex.parallel import DistributedDataParallel
+from dllogger import JSONStreamBackend, StdOutBackend, Verbosity
+
+from jasper import config
+from common import helpers
+from common.dali.data_loader import DaliDataLoader
+from common.dataset import (AudioDataset, FilelistDataset, get_data_loader,
+                            SingleAudioDataset)
+from common.features import BaseFeatures, FilterbankFeatures
+from common.helpers import print_once, process_evaluation_epoch
+from jasper.model import GreedyCTCDecoder, Jasper
+from common.tb_dllogger import stdout_metric_format, unique_log_fpath
+
+
+def get_parser():
+    parser = argparse.ArgumentParser(description='Jasper')
+    parser.add_argument('--batch_size', default=16, type=int,
+                        help='Data batch size')
+    parser.add_argument('--steps', default=0, type=int,
+                        help='Eval this many steps for every worker')
+    parser.add_argument('--warmup_steps', default=0, type=int,
+                        help='Burn-in period before measuring latencies')
+    parser.add_argument('--model_config', type=str, required=True,
+                        help='Relative model config path given dataset folder')
+    parser.add_argument('--dataset_dir', type=str,
+                        help='Absolute path to dataset folder')
+    parser.add_argument('--val_manifests', type=str, nargs='+',
+                        help='Relative path to evaluation dataset manifest files')
+    parser.add_argument('--ckpt', default=None, type=str,
+                        help='Path to model checkpoint')
+    parser.add_argument('--pad_leading', type=int, default=16,
+                        help='Pads every batch with leading zeros '
+                             'to counteract conv shifts of the field of view')
+    parser.add_argument('--amp', '--fp16', action='store_true',
+                        help='Use FP16 precision')
+    parser.add_argument('--cudnn_benchmark', action='store_true',
+                        help='Enable cudnn benchmark')
+    parser.add_argument('--cpu', action='store_true',
+                        help='Run inference on CPU')
+    parser.add_argument("--seed", default=None, type=int, help='Random seed')
+    parser.add_argument('--local_rank', default=os.getenv('LOCAL_RANK', 0),
+                        type=int, help='GPU id used for distributed training')
+
+    io = parser.add_argument_group('feature and checkpointing setup')
+    io.add_argument('--dali_device', type=str, choices=['none', 'cpu', 'gpu'],
+                    default='gpu', help='Use DALI pipeline for fast data processing')
+    io.add_argument('--save_predictions', type=str, default=None,
+                    help='Save predictions in text form at this location')
+    io.add_argument('--save_logits', default=None, type=str,
+                    help='Save output logits under specified path')
+    io.add_argument('--transcribe_wav', type=str,
+                    help='Path to a single .wav file (16KHz)')
+    io.add_argument('--transcribe_filelist', type=str,
+                    help='Path to a filelist with one .wav path per line')
+    io.add_argument('-o', '--output_dir', default='results/',
+                    help='Output folder to save audio (file per phrase)')
+    io.add_argument('--log_file', type=str, default=None,
+                    help='Path to a DLLogger log file')
+    io.add_argument('--ema', action='store_true',
+                    help='Load averaged model weights')
+    io.add_argument('--torchscript', action='store_true',
+                    help='Evaluate with a TorchScripted model')
+    io.add_argument('--torchscript_export', action='store_true',
+                    help='Export the model with torch.jit to the output_dir')
+    io.add_argument('--override_config', type=str, action='append',
+                    help='Overrides a value from a config .yaml.'
+                         ' Syntax: `--override_config nested.config.key=val`.')
+    return parser
+
+
+def durs_to_percentiles(durations, ratios):
+    durations = np.asarray(durations) * 1000  # in ms
+    latency = durations
+
+    latency = latency[5:]
+    mean_latency = np.mean(latency)
+
+    latency_worst = nlargest(math.ceil((1 - min(ratios)) * len(latency)), latency)
+    latency_ranges = get_percentile(ratios, latency_worst, len(latency))
+    latency_ranges[0.5] = mean_latency
+    return latency_ranges
+
+
+def get_percentile(ratios, arr, nsamples):
+    res = {}
+    for a in ratios:
+        idx = max(int(nsamples * (1 - a)), 0)
+        res[a] = arr[idx]
+    return res
+
+
+def torchscript_export(data_loader, audio_processor, model, greedy_decoder,
+                       output_dir, use_amp, use_conv_masks, model_config, device,
+                       save):
+
+    audio_processor.to(device)
+
+    for batch in data_loader:
+        batch = [t.to(device, non_blocking=True) for t in batch]
+        audio, audio_len, _, _ = batch
+        feats, feat_lens = audio_processor(audio, audio_len)
+        break
+
+    print("\nExporting featurizer...")
+    print("\nNOTE: Dithering causes warnings about non-determinism.\n")
+    ts_feat = torch.jit.trace(audio_processor, (audio, audio_len))
+
+    print("\nExporting acoustic model...")
+    model(feats, feat_lens)
+    ts_acoustic = torch.jit.trace(model, (feats, feat_lens))
+
+    print("\nExporting decoder...")
+    log_probs = model(feats, feat_lens)
+    ts_decoder = torch.jit.script(greedy_decoder, log_probs)
+    print("\nJIT export complete.")
+
+    if save:
+        precision = "fp16" if use_amp else "fp32"
+        module_name = f'{os.path.basename(model_config)}_{precision}'
+        ts_feat.save(os.path.join(output_dir, module_name + "_feat.pt"))
+        ts_acoustic.save(os.path.join(output_dir, module_name + "_acoustic.pt"))
+        ts_decoder.save(os.path.join(output_dir, module_name + "_decoder.pt"))
+
+    return ts_feat, ts_acoustic, ts_decoder
+
+
+def main():
+
+    parser = get_parser()
+    args = parser.parse_args()
+
+    log_fpath = args.log_file or str(Path(args.output_dir, 'nvlog_infer.json'))
+    log_fpath = unique_log_fpath(log_fpath)
+    dllogger.init(backends=[JSONStreamBackend(Verbosity.DEFAULT, log_fpath),
+                            StdOutBackend(Verbosity.VERBOSE,
+                                          metric_format=stdout_metric_format)])
+
+    [dllogger.log("PARAMETER", {k: v}) for k, v in vars(args).items()]
+
+    for step in ['DNN', 'data+DNN', 'data']:
+        for c in [0.99, 0.95, 0.9, 0.5]:
+            cs = 'avg' if c == 0.5 else f'{int(100*c)}%'
+            dllogger.metadata(f'{step.lower()}_latency_{c}',
+                              {'name': f'{step} latency {cs}',
+                               'format': ':>7.2f', 'unit': 'ms'})
+    dllogger.metadata(
+        'eval_wer', {'name': 'WER', 'format': ':>3.2f', 'unit': '%'})
+
+    if args.cpu:
+        device = torch.device('cpu')
+    else:
+        assert torch.cuda.is_available()
+        device = torch.device('cuda')
+        torch.backends.cudnn.benchmark = args.cudnn_benchmark
+
+    if args.seed is not None:
+        torch.manual_seed(args.seed + args.local_rank)
+        np.random.seed(args.seed + args.local_rank)
+        random.seed(args.seed + args.local_rank)
+
+    # set up distributed training
+    multi_gpu = not args.cpu and int(os.environ.get('WORLD_SIZE', 1)) > 1
+    if multi_gpu:
+        torch.cuda.set_device(args.local_rank)
+        distrib.init_process_group(backend='nccl', init_method='env://')
+        print_once(f'Inference with {distrib.get_world_size()} GPUs')
+
+    cfg = config.load(args.model_config)
+    config.apply_config_overrides(cfg, args)
+
+    symbols = helpers.add_ctc_blank(cfg['labels'])
+
+    use_dali = args.dali_device in ('cpu', 'gpu')
+    dataset_kw, features_kw = config.input(cfg, 'val')
+
+    measure_perf = args.steps > 0
+
+    # dataset
+    if args.transcribe_wav or args.transcribe_filelist:
+
+        if use_dali:
+            print("DALI supported only with input .json files; disabling")
+            use_dali = False
+
+        assert not (args.transcribe_wav and args.transcribe_filelist)
+
+        if args.transcribe_wav:
+            dataset = SingleAudioDataset(args.transcribe_wav)
+        else:
+            dataset = FilelistDataset(args.transcribe_filelist)
+
+        data_loader = get_data_loader(dataset,
+                                      batch_size=1,
+                                      multi_gpu=multi_gpu,
+                                      shuffle=False,
+                                      num_workers=0,
+                                      drop_last=(True if measure_perf else False))
+
+        _, features_kw = config.input(cfg, 'val')
+        assert not features_kw['pad_to_max_duration']
+        feat_proc = FilterbankFeatures(**features_kw)
+
+    elif use_dali:
+        # pad_to_max_duration is not supported by DALI - have simple padders
+        if features_kw['pad_to_max_duration']:
+            feat_proc = BaseFeatures(
+                pad_align=features_kw['pad_align'],
+                pad_to_max_duration=True,
+                max_duration=features_kw['max_duration'],
+                sample_rate=features_kw['sample_rate'],
+                window_size=features_kw['window_size'],
+                window_stride=features_kw['window_stride'])
+            features_kw['pad_to_max_duration'] = False
+        else:
+            feat_proc = None
+
+        data_loader = DaliDataLoader(
+            gpu_id=args.local_rank or 0,
+            dataset_path=args.dataset_dir,
+            config_data=dataset_kw,
+            config_features=features_kw,
+            json_names=args.val_manifests,
+            batch_size=args.batch_size,
+            pipeline_type=("train" if measure_perf else "val"),  # no drop_last
+            device_type=args.dali_device,
+            symbols=symbols)
+
+    else:
+        dataset = AudioDataset(args.dataset_dir,
+                               args.val_manifests,
+                               symbols,
+                               **dataset_kw)
+
+        data_loader = get_data_loader(dataset,
+                                      args.batch_size,
+                                      multi_gpu=multi_gpu,
+                                      shuffle=False,
+                                      num_workers=4,
+                                      drop_last=False)
+
+        feat_proc = FilterbankFeatures(**features_kw)
+
+    model = Jasper(encoder_kw=config.encoder(cfg),
+                   decoder_kw=config.decoder(cfg, n_classes=len(symbols)))
+
+    if args.ckpt is not None:
+        print(f'Loading the model from {args.ckpt} ...')
+        checkpoint = torch.load(args.ckpt, map_location="cpu")
+        key = 'ema_state_dict' if args.ema else 'state_dict'
+        state_dict = helpers.convert_v1_state_dict(checkpoint[key])
+        model.load_state_dict(state_dict, strict=True)
+
+    model.to(device)
+    model.eval()
+
+    if feat_proc is not None:
+        feat_proc.to(device)
+        feat_proc.eval()
+
+    if args.amp:
+        model = model.half()
+
+    if args.torchscript:
+        greedy_decoder = GreedyCTCDecoder()
+
+        feat_proc, model, greedy_decoder = torchscript_export(
+            data_loader, feat_proc, model, greedy_decoder, args.output_dir,
+            use_amp=args.amp, use_conv_masks=True, model_toml=args.model_toml,
+            device=device, save=args.torchscript_export)
+
+    if multi_gpu:
+        model = DistributedDataParallel(model)
+
+    agg = {'txts': [], 'preds': [], 'logits': []}
+    dur = {'data': [], 'dnn': [], 'data+dnn': []}
+
+    looped_loader = chain.from_iterable(repeat(data_loader))
+    greedy_decoder = GreedyCTCDecoder()
+
+    sync = lambda: torch.cuda.synchronize() if device.type == 'cuda' else None
+
+    steps = args.steps + args.warmup_steps or len(data_loader)
+    with torch.no_grad():
+
+        for it, batch in enumerate(tqdm(looped_loader, initial=1, total=steps)):
+
+            if use_dali:
+                feats, feat_lens, txt, txt_lens = batch
+                if feat_proc is not None:
+                    feats, feat_lens = feat_proc(feats, feat_lens)
+            else:
+                batch = [t.to(device, non_blocking=True) for t in batch]
+                audio, audio_lens, txt, txt_lens = batch
+                feats, feat_lens = feat_proc(audio, audio_lens)
+
+            sync()
+            t1 = time.perf_counter()
+
+            if args.amp:
+                feats = feats.half()
+
+            feats = F.pad(feats, (args.pad_leading, 0))
+            feat_lens += args.pad_leading
+
+            if model.encoder.use_conv_masks:
+                log_probs, log_prob_lens = model(feats, feat_lens)
+            else:
+                log_probs = model(feats, feat_lens)
+
+            preds = greedy_decoder(log_probs)
+
+            sync()
+            t2 = time.perf_counter()
+
+            # burn-in period; wait for a new loader due to num_workers
+            if it >= 1 and (args.steps == 0 or it >= args.warmup_steps):
+                dur['data'].append(t1 - t0)
+                dur['dnn'].append(t2 - t1)
+                dur['data+dnn'].append(t2 - t0)
+
+            if txt is not None:
+                agg['txts'] += helpers.gather_transcripts([txt], [txt_lens],
+                                                          symbols)
+            agg['preds'] += helpers.gather_predictions([preds], symbols)
+            agg['logits'].append(log_probs)
+
+            if it + 1 == steps:
+                break
+
+            sync()
+            t0 = time.perf_counter()
+
+        # communicate the results
+        if args.transcribe_wav:
+            for idx, p in enumerate(agg['preds']):
+                print_once(f'Prediction {idx+1: >3}: {p}')
+
+        elif args.transcribe_filelist:
+            pass
+
+        elif not multi_gpu or distrib.get_rank() == 0:
+            wer, _ = process_evaluation_epoch(agg)
+
+            dllogger.log(step=(), data={'eval_wer': 100 * wer})
+
+        if args.save_predictions:
+            with open(args.save_predictions, 'w') as f:
+                f.write('\n'.join(agg['preds']))
+
+        if args.save_logits:
+            logits = torch.cat(agg['logits'], dim=0).cpu()
+            torch.save(logits, args.save_logits)
+
+    # report timings
+    if len(dur['data']) >= 20:
+        ratios = [0.9, 0.95, 0.99]
+        for stage in dur:
+            lat = durs_to_percentiles(dur[stage], ratios)
+            for k in [0.99, 0.95, 0.9, 0.5]:
+                kk = str(k).replace('.', '_')
+                dllogger.log(step=(), data={f'{stage.lower()}_latency_{kk}': lat[k]})
+
+    else:
+        print_once('Not enough samples to measure latencies.')
+
+
+if __name__ == "__main__":
+    main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/jasper/config.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/jasper/config.py
new file mode 100644
index 0000000000..60f7a258eb
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/jasper/config.py
@@ -0,0 +1,125 @@
+import copy
+import inspect
+import typing
+from ast import literal_eval
+from contextlib import suppress
+from numbers import Number
+
+import yaml
+
+from .model import JasperDecoderForCTC, JasperBlock, JasperEncoder
+from common.audio import GainPerturbation, ShiftPerturbation, SpeedPerturbation
+from common.dataset import AudioDataset
+from common.features import CutoutAugment, FilterbankFeatures, SpecAugment
+from common.helpers import print_once
+
+
+def default_args(klass):
+    sig = inspect.signature(klass.__init__)
+    return {k: v.default for k,v in sig.parameters.items() if k != 'self'}
+
+
+def load(fpath):
+    if fpath.endswith('.toml'):
+        raise ValueError('.toml config format has been changed to .yaml')
+
+    cfg = yaml.safe_load(open(fpath, 'r'))
+
+    # Reload to deep copy shallow copies, which were made with yaml anchors
+    yaml.Dumper.ignore_aliases = lambda *args: True
+    cfg = yaml.dump(cfg)
+    cfg = yaml.safe_load(cfg)
+    return cfg
+
+
+def validate_and_fill(klass, user_conf, ignore_unk=[], optional=[]):
+    conf = default_args(klass)
+
+    for k,v in user_conf.items():
+        assert k in conf or k in ignore_unk, f'Unknown parameter {k} for {klass}'
+        conf[k] = v
+
+    # Keep only mandatory or optional-nonempty
+    conf = {k:v for k,v in conf.items()
+            if k not in optional or v is not inspect.Parameter.empty}
+
+    # Validate
+    for k,v in conf.items():
+        assert v is not inspect.Parameter.empty, \
+            f'Value for {k} not specified for {klass}'
+    return conf
+
+
+def input(conf_yaml, split='train'):
+    conf = copy.deepcopy(conf_yaml[f'input_{split}'])
+    conf_dataset = conf.pop('audio_dataset')
+    conf_features = conf.pop('filterbank_features')
+
+    # Validate known inner classes
+    inner_classes = [
+        (conf_dataset, 'speed_perturbation', SpeedPerturbation),
+        (conf_dataset, 'gain_perturbation', GainPerturbation),
+        (conf_dataset, 'shift_perturbation', ShiftPerturbation),
+        (conf_features, 'spec_augment', SpecAugment),
+        (conf_features, 'cutout_augment', CutoutAugment),
+    ]
+    for conf_tgt, key, klass in inner_classes:
+        if key in conf_tgt:
+            conf_tgt[key] = validate_and_fill(klass, conf_tgt[key])
+
+    for k in conf:
+        raise ValueError(f'Unknown key {k}')
+
+    # Validate outer classes
+    conf_dataset = validate_and_fill(
+        AudioDataset, conf_dataset,
+        optional=['data_dir', 'labels', 'manifest_fpaths'])
+
+    conf_features = validate_and_fill(
+        FilterbankFeatures, conf_features)
+
+    # Check params shared between classes
+    shared = ['sample_rate', 'max_duration', 'pad_to_max_duration']
+    for sh in shared:
+        assert conf_dataset[sh] == conf_features[sh], (
+            f'{sh} should match in Dataset and FeatureProcessor: '
+            f'{conf_dataset[sh]}, {conf_features[sh]}')
+
+    return conf_dataset, conf_features
+
+
+def encoder(conf):
+    """Validate config for JasperEncoder and subsequent JasperBlocks"""
+
+    # Validate, but don't overwrite with defaults
+    for blk in conf['jasper']['encoder']['blocks']:
+        validate_and_fill(JasperBlock, blk, optional=['infilters'],
+                          ignore_unk=['residual_dense'])
+
+    return validate_and_fill(JasperEncoder, conf['jasper']['encoder'])
+
+
+def decoder(conf, n_classes):
+    decoder_kw = {'n_classes': n_classes, **conf['jasper']['decoder']}
+    return validate_and_fill(JasperDecoderForCTC, decoder_kw)
+
+
+def apply_config_overrides(conf, args):
+    if args.override_config is None:
+        return
+    for override_key_val in args.override_config:
+        key, val = override_key_val.split('=')
+        with suppress(TypeError, ValueError):
+            val = literal_eval(val)
+        apply_nested_config_override(conf, key, val)
+
+
+def apply_nested_config_override(conf, key_str, val):
+    fields = key_str.split('.')
+    for f in fields[:-1]:
+        conf = conf[f]
+    f = fields[-1]
+    assert (f not in conf
+            or type(val) is type(conf[f])
+            or (isinstance(val, Number) and isinstance(conf[f], Number)))
+    conf[f] = val
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/jasper/model.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/jasper/model.py
new file mode 100644
index 0000000000..dd38ce4b77
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/jasper/model.py
@@ -0,0 +1,275 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import torch
+import torch.nn as nn
+import torch.nn.functional as F
+
+
+activations = {
+    "hardtanh": nn.Hardtanh,
+    "relu": nn.ReLU,
+    "selu": nn.SELU,
+}
+
+
+def init_weights(m, mode='xavier_uniform'):
+    if type(m) == nn.Conv1d or type(m) == MaskedConv1d:
+        if mode == 'xavier_uniform':
+            nn.init.xavier_uniform_(m.weight, gain=1.0)
+        elif mode == 'xavier_normal':
+            nn.init.xavier_normal_(m.weight, gain=1.0)
+        elif mode == 'kaiming_uniform':
+            nn.init.kaiming_uniform_(m.weight, nonlinearity="relu")
+        elif mode == 'kaiming_normal':
+            nn.init.kaiming_normal_(m.weight, nonlinearity="relu")
+        else:
+            raise ValueError("Unknown Initialization mode: {0}".format(mode))
+
+    elif type(m) == nn.BatchNorm1d:
+        if m.track_running_stats:
+            m.running_mean.zero_()
+            m.running_var.fill_(1)
+            m.num_batches_tracked.zero_()
+        if m.affine:
+            nn.init.ones_(m.weight)
+            nn.init.zeros_(m.bias)
+
+
+def get_same_padding(kernel_size, stride, dilation):
+    if stride > 1 and dilation > 1:
+        raise ValueError("Only stride OR dilation may be greater than 1")
+    return (kernel_size // 2) * dilation
+
+
+class MaskedConv1d(nn.Conv1d):
+    """1D convolution with sequence masking
+    """
+    __constants__ = ["masked"]
+    def __init__(self, in_channels, out_channels, kernel_size, stride=1,
+                 padding=0, dilation=1, groups=1, bias=False, masked=True):
+        super(MaskedConv1d, self).__init__(
+            in_channels, out_channels, kernel_size, stride=stride,
+            padding=padding, dilation=dilation, groups=groups, bias=bias)
+
+        self.masked = masked
+
+    def get_seq_len(self, lens):
+        return ((lens + 2 * self.padding[0] - self.dilation[0]
+                 * (self.kernel_size[0] - 1) - 1) // self.stride[0] + 1)
+
+    def forward(self, x, x_lens=None):
+        if self.masked:
+            max_len = x.size(2)
+            idxs = torch.arange(max_len, dtype=x_lens.dtype, device=x_lens.device)
+            mask = idxs.expand(x_lens.size(0), max_len) >= x_lens.unsqueeze(1)
+            x = x.masked_fill(mask.unsqueeze(1).to(device=x.device), 0)
+            x_lens = self.get_seq_len(x_lens)
+
+        return super(MaskedConv1d, self).forward(x), x_lens
+
+
+class JasperBlock(nn.Module):
+    __constants__ = ["use_conv_masks"]
+
+    """Jasper Block. See https://arxiv.org/pdf/1904.03288.pdf
+    """
+    def __init__(self, infilters, filters, repeat=3, kernel_size=11, stride=1,
+                 dilation=1, padding='same', dropout=0.2, activation=None,
+                 residual=True, residual_panes=[], use_conv_masks=False):
+        super(JasperBlock, self).__init__()
+
+        assert padding == "same", "Only 'same' padding is supported."
+
+        padding_val = get_same_padding(kernel_size[0], stride[0], dilation[0])
+        self.use_conv_masks = use_conv_masks
+        self.conv = nn.ModuleList()
+        for i in range(repeat):
+            self.conv.extend(self._conv_bn(infilters if i == 0 else filters,
+                                           filters,
+                                           kernel_size=kernel_size,
+                                           stride=stride,
+                                           dilation=dilation,
+                                           padding=padding_val))
+            if i < repeat - 1:
+                self.conv.extend(self._act_dropout(dropout, activation))
+
+        self.res = nn.ModuleList() if residual else None
+        res_panes = residual_panes.copy()
+        self.dense_residual = residual
+
+        if residual:
+            if len(residual_panes) == 0:
+                res_panes = [infilters]
+                self.dense_residual = False
+
+            for ip in res_panes:
+                self.res.append(nn.ModuleList(
+                    self._conv_bn(ip, filters, kernel_size=1)))
+
+        self.out = nn.Sequential(*self._act_dropout(dropout, activation))
+
+    def _conv_bn(self, in_channels, out_channels, **kw):
+        return [MaskedConv1d(in_channels, out_channels,
+                             masked=self.use_conv_masks, **kw),
+                nn.BatchNorm1d(out_channels, eps=1e-3, momentum=0.1)]
+
+    def _act_dropout(self, dropout=0.2, activation=None):
+        return [activation or nn.Hardtanh(min_val=0.0, max_val=20.0),
+                nn.Dropout(p=dropout)]
+
+    def forward(self, xs, xs_lens=None):
+        if not self.use_conv_masks:
+            xs_lens = 0
+
+        # forward convolutions
+        out = xs[-1]
+        lens = xs_lens
+        for i, l in enumerate(self.conv):
+            if isinstance(l, MaskedConv1d):
+                out, lens = l(out, lens)
+            else:
+                out = l(out)
+
+        # residuals
+        if self.res is not None:
+            for i, layer in enumerate(self.res):
+                res_out = xs[i]
+                for j, res_layer in enumerate(layer):
+                    if j == 0:  #  and self.use_conv_mask:
+                        res_out, _ = res_layer(res_out, xs_lens)
+                    else:
+                        res_out = res_layer(res_out)
+                out += res_out
+
+        # output
+        out = self.out(out)
+        if self.res is not None and self.dense_residual:
+            out = xs + [out]
+        else:
+            out = [out]
+
+        if self.use_conv_masks:
+            return out, lens
+        else:
+            return out, None
+
+
+class JasperEncoder(nn.Module):
+    __constants__ = ["use_conv_masks"]
+
+    def __init__(self, in_feats, activation, frame_splicing=1,
+                 init='xavier_uniform', use_conv_masks=False, blocks=[]):
+        super(JasperEncoder, self).__init__()
+
+        self.use_conv_masks = use_conv_masks
+        self.layers = nn.ModuleList()
+
+        in_feats *= frame_splicing
+        all_residual_panes = []
+        for i,blk in enumerate(blocks):
+
+            blk['activation'] = activations[activation]()
+
+            has_residual_dense = blk.pop('residual_dense', False)
+            if has_residual_dense:
+                all_residual_panes += [in_feats]
+                blk['residual_panes'] = all_residual_panes
+            else:
+                blk['residual_panes'] = []
+
+            self.layers.append(
+                JasperBlock(in_feats, use_conv_masks=use_conv_masks, **blk))
+
+            in_feats = blk['filters']
+
+        self.apply(lambda x: init_weights(x, mode=init))
+
+    def forward(self, x, x_lens=None):
+        out, out_lens = [x], x_lens
+        for l in self.layers:
+            out, out_lens = l(out, out_lens)
+
+        return out, out_lens
+
+
+class JasperDecoderForCTC(nn.Module):
+    def __init__(self, in_feats, n_classes, init='xavier_uniform'):
+        super(JasperDecoderForCTC, self).__init__()
+
+        self.layers = nn.Sequential(
+            nn.Conv1d(in_feats, n_classes, kernel_size=1, bias=True),)
+        self.apply(lambda x: init_weights(x, mode=init))
+
+    def forward(self, enc_out):
+        out = self.layers(enc_out[-1]).transpose(1, 2)
+        return F.log_softmax(out, dim=2)
+
+
+class GreedyCTCDecoder(nn.Module):
+    @torch.no_grad()
+    def forward(self, log_probs, log_prob_lens=None):
+
+        if log_prob_lens is not None:
+            max_len = log_probs.size(1)
+            idxs = torch.arange(max_len, dtype=log_prob_lens.dtype,
+                                device=log_prob_lens.device)
+            mask = idxs.unsqueeze(0) >= log_prob_lens.unsqueeze(1)
+            log_probs[:,:,-1] = log_probs[:,:,-1].masked_fill(mask, float("Inf"))
+
+        return log_probs.argmax(dim=-1, keepdim=False).int()
+
+
+class Jasper(nn.Module):
+    def __init__(self, encoder_kw, decoder_kw, transpose_in=False):
+        super(Jasper, self).__init__()
+        self.transpose_in = transpose_in
+        self.encoder = JasperEncoder(**encoder_kw)
+        self.decoder = JasperDecoderForCTC(**decoder_kw)
+
+    def forward(self, x, x_lens=None):
+        if self.encoder.use_conv_masks:
+            assert x_lens is not None
+            enc, enc_lens = self.encoder(x, x_lens)
+            out = self.decoder(enc)
+            return out, enc_lens
+        else:
+            if self.transpose_in:
+                x = x.transpose(1, 2)
+            enc, _ = self.encoder(x)
+            out = self.decoder(enc)
+            return out  # torchscript refuses to output None
+
+    # TODO Explicitly add x_lens=None for inference (now x can be a Tensor or tuple)
+    def infer(self, x, x_lens=None):
+        if self.encoder.use_conv_masks:
+            return self.forward(x, x_lens)
+        else:
+            ret = self.forward(x)
+            return ret, len(ret)
+
+
+class CTCLossNM:
+    def __init__(self, n_classes):
+        self._criterion = nn.CTCLoss(blank=n_classes-1, reduction='none')
+
+    def __call__(self, log_probs, targets, input_length, target_length):
+        input_length = input_length.long()
+        target_length = target_length.long()
+        targets = targets.long()
+        loss = self._criterion(log_probs.transpose(1, 0), targets, input_length,
+                               target_length)
+        # note that this is different from reduction = 'mean'
+        # because we are not dividing by target lengths
+        return torch.mean(loss)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/notebooks/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/notebooks/README.md
new file mode 100644
index 0000000000..072ed740fa
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/notebooks/README.md
@@ -0,0 +1,203 @@
+# Jasper notebooks
+
+This folder provides different notebooks to run Jasper inference step by step. 
+
+## Table Of Contents
+
+- [Jasper Jupyter Notebook for TensorRT](#jasper-jupyter-notebook-for-tensorrt)
+   * [Requirements](#requirements)
+    * [Quick Start Guide](#quick-start-guide)
+- [Jasper Colab Notebook for TensorRT](#jasper-colab-notebook-for-tensorrt)
+   * [Requirements](#requirements)
+    * [Quick Start Guide](#quick-start-guide)
+- [Jasper Jupyter Notebook for TensorRT Inference Server](#jasper-colab-notebook-for-tensorrt-inference-server)
+   * [Requirements](#requirements)
+    * [Quick Start Guide](#quick-start-guide)
+
+## Jasper Jupyter Notebook for TensorRT
+### Requirements
+
+`./trt/` contains a Dockerfile which extends the PyTorch 19.09-py3 NGC container and encapsulates some dependencies. Aside from these dependencies, ensure you have the following components:
+
+* [NVIDIA Turing](https://www.nvidia.com/en-us/geforce/turing/) or [Volta](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/) based GPU    
+* [NVIDIA Docker](https://github.com/NVIDIA/nvidia-docker)
+* [PyTorch 19.09-py3 NGC container](https://ngc.nvidia.com/catalog/containers/nvidia:pytorch)
+* [NVIDIA machine learning repository](https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb) and [NVIDIA cuda repository](https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-repo-ubuntu1804_10.1.243-1_amd64.deb) for NVIDIA TensorRT 6
+* [NVIDIA Volta](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/) or [Turing](https://www.nvidia.com/en-us/geforce/turing/) based GPU
+* [Pretrained Jasper Model Checkpoint](https://ngc.nvidia.com/catalog/models/nvidia:jasperpyt_fp16)
+
+### Quick Start Guide
+
+Running the following scripts will build and launch the container containing all required dependencies for both TensorRT as well as native PyTorch. This is necessary for using inference with TensorRT and can also be used for data download, processing and training of the model.
+
+#### 1. Clone the repository.
+
+```
+git clone https://github.com/NVIDIA/DeepLearningExamples
+cd DeepLearningExamples/PyTorch/SpeechRecognition/Jasper
+```
+
+#### 2. Build the Jasper PyTorch with TRT 6 container:
+
+```
+bash trt/scripts/docker/build.sh
+```
+
+#### 3. Create directories
+Prepare to start a detached session in the NGC container.
+Create three directories on your local machine for dataset, checkpoint, and result, respectively, naming "data" "checkpoint" "result":
+
+```
+mkdir data checkpoint result
+```
+
+#### 4. Download the checkpoint
+Download the checkpoint file jasperpyt_fp16 from NGC Model Repository:  
+- https://ngc.nvidia.com/catalog/models/nvidia:jasperpyt_fp16
+
+to the directory: _checkpoint_
+
+The Jasper PyTorch container will be launched in the Jupyter notebook. Within the container, the contents of the root repository will be copied to the /workspace/jasper directory.
+
+The /datasets, /checkpoints, /results directories are mounted as volumes and mapped to the corresponding directories "data" "checkpoint" "result" on the host.
+
+#### 5. Run the notebook
+
+For running the notebook on your local machine, run:
+
+```
+jupyter notebook -- notebooks/JasperTRT.ipynb
+```
+
+For running the notebook on another machine remotely, run:
+
+```
+jupyter notebook --ip=0.0.0.0 --allow-root
+```
+
+And navigate a web browser to the IP address or hostname of the host machine at port 8888: `http://[host machine]:8888`
+
+Use the token listed in the output from running the jupyter command to log in, for example: `http://[host machine]:8888/?token=aae96ae9387cd28151868fee318c3b3581a2d794f3b25c6b`
+
+
+
+## Jasper Colab Notebook for TensorRT
+### Requirements
+
+`./trt/` contains a Dockerfile which extends the PyTorch 19.09-py3 NGC container and encapsulates some dependencies. Aside from these dependencies, ensure you have the following components:
+
+* [NVIDIA Turing](https://www.nvidia.com/en-us/geforce/turing/) or [Volta](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/) based GPU    
+* [NVIDIA Docker](https://github.com/NVIDIA/nvidia-docker)
+* [PyTorch 19.09-py3 NGC container](https://ngc.nvidia.com/catalog/containers/nvidia:pytorch)
+* [NVIDIA machine learning repository](https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb) and [NVIDIA cuda repository](https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-repo-ubuntu1804_10.1.243-1_amd64.deb) for NVIDIA TensorRT 6
+* [NVIDIA Volta](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/) or [Turing](https://www.nvidia.com/en-us/geforce/turing/) based GPU
+* [Pretrained Jasper Model Checkpoint](https://ngc.nvidia.com/catalog/models/nvidia:jasperpyt_fp16)
+
+### Quick Start Guide
+
+Running the following scripts will build and launch the container containing all required dependencies for both TensorRT as well as native PyTorch. This is necessary for using inference with TensorRT and can also be used for data download, processing and training of the model.
+
+#### 1. Clone the repository.
+
+```
+git clone https://github.com/NVIDIA/DeepLearningExamples
+cd DeepLearningExamples/PyTorch/SpeechRecognition/Jasper
+```
+
+#### 2. Build the Jasper PyTorch with TRT 6 container:
+
+```
+bash trt/scripts/docker/build.sh
+```
+
+#### 3. Create directories
+Prepare to start a detached session in the NGC container.
+Create three directories on your local machine for dataset, checkpoint, and result, respectively, naming "data" "checkpoint" "result":
+
+```
+mkdir data checkpoint result
+```
+
+#### 4. Download the checkpoint
+Download the checkpoint file jasperpyt_fp16 from NGC Model Repository:  
+- https://ngc.nvidia.com/catalog/models/nvidia:jasperpyt_fp16
+
+to the directory: _checkpoint_
+
+The Jasper PyTorch container will be launched in the Jupyter notebook. Within the container, the contents of the root repository will be copied to the /workspace/jasper directory.
+
+The /datasets, /checkpoints, /results directories are mounted as volumes and mapped to the corresponding directories "data" "checkpoint" "result" on the host.
+
+#### 5. Run the notebook
+
+>>>>>>> 2deaddbc2ea58d5318b06203ae30ace2dd576ecb
+For running the notebook on your local machine, run:
+
+```
+jupyter notebook -- notebooks/Colab_Jasper_TRT_inference_demo.ipynb
+```
+
+For running the notebook on another machine remotely, run:
+
+```
+jupyter notebook --ip=0.0.0.0 --allow-root
+```
+
+And navigate a web browser to the IP address or hostname of the host machine at port 8888: `http://[host machine]:8888`
+
+Use the token listed in the output from running the jupyter command to log in, for example: `http://[host machine]:8888/?token=aae96ae9387cd28151868fee318c3b3581a2d794f3b25c6b`
+
+
+
+## Jasper Jupyter Notebook for TensorRT Inference Server
+### Requirements
+
+`./trtis/` contains a Dockerfile which extends the PyTorch 19.09-py3 NGC container and encapsulates some dependencies. Aside from these dependencies, ensure you have the following components:
+
+* [NVIDIA Turing](https://www.nvidia.com/en-us/geforce/turing/) or [Volta](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/) based GPU    
+* [NVIDIA Docker](https://github.com/NVIDIA/nvidia-docker)
+* [PyTorch 19.09-py3 NGC container](https://ngc.nvidia.com/catalog/containers/nvidia:pytorch)
+* [TensorRT Inference Server 19.09 NGC container](https://ngc.nvidia.com/catalog/containers/nvidia:tensorrtserver)
+* [NVIDIA machine learning repository](https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb) and [NVIDIA cuda repository](https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-repo-ubuntu1804_10.1.243-1_amd64.deb) for NVIDIA TensorRT 6
+* [NVIDIA Volta](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/) or [Turing](https://www.nvidia.com/en-us/geforce/turing/) based GPU
+* [Pretrained Jasper Model Checkpoint](https://ngc.nvidia.com/catalog/models/nvidia:jasperpyt_fp16)
+
+### Quick Start Guide
+
+
+#### 1. Clone the repository.
+
+```
+git clone https://github.com/NVIDIA/DeepLearningExamples
+cd DeepLearningExamples/PyTorch/SpeechRecognition/Jasper
+```
+
+#### 2. Build a container that extends NGC PyTorch 19.09, TensorRT, TensorRT Inference Server, and TensorRT Inference Client.
+
+```
+bash trtis/scripts/docker/build.sh
+```
+
+#### 3. Download the checkpoint
+Download the checkpoint file jasper_fp16.pt from NGC Model Repository:  
+- https://ngc.nvidia.com/catalog/models/nvidia:jasperpyt_fp16
+
+to an user specified directory _CHECKPOINT_DIR_
+
+#### 4. Run the notebook
+
+For running the notebook on your local machine, run:
+
+```
+jupyter notebook -- notebooks/JasperTRTIS.ipynb
+```
+
+For running the notebook on another machine remotely, run:
+
+```
+jupyter notebook --ip=0.0.0.0 --allow-root
+```
+
+And navigate a web browser to the IP address or hostname of the host machine at port 8888: `http://[host machine]:8888`
+
+Use the token listed in the output from running the jupyter command to log in, for example: `http://[host machine]:8888/?token=aae96ae9387cd28151868fee318c3b3581a2d794f3b25c6b`
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-16GB_Jasper_AMP_8GPU.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-16GB_Jasper_AMP_8GPU.sh
new file mode 100644
index 0000000000..57bcd4c5ec
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-16GB_Jasper_AMP_8GPU.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+NUM_GPUS=8 AMP=true BATCH_SIZE=64 GRADIENT_ACCUMULATION_STEPS=4 bash scripts/train.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-16GB_Jasper_FP32_8GPU.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-16GB_Jasper_FP32_8GPU.sh
new file mode 100644
index 0000000000..0317b41a61
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-16GB_Jasper_FP32_8GPU.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+NUM_GPUS=8 BATCH_SIZE=64 GRADIENT_ACCUMULATION_STEPS=4 bash scripts/train.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-32GB_Jasper_AMP_8GPU.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-32GB_Jasper_AMP_8GPU.sh
new file mode 100644
index 0000000000..8953566d4a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-32GB_Jasper_AMP_8GPU.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+NUM_GPUS=8 AMP=true BATCH_SIZE=64 GRADIENT_ACCUMULATION_STEPS=1 bash scripts/train.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-32GB_Jasper_FP32_8GPU.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-32GB_Jasper_FP32_8GPU.sh
new file mode 100644
index 0000000000..eed6a1273f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX1-32GB_Jasper_FP32_8GPU.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+NUM_GPUS=8 BATCH_SIZE=64 GRADIENT_ACCUMULATION_STEPS=2 bash scripts/train.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_AMP_16GPU.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_AMP_16GPU.sh
new file mode 100644
index 0000000000..a0ccca4a1d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_AMP_16GPU.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+NUM_GPUS=16 AMP=true BATCH_SIZE=64 GRADIENT_ACCUMULATION_STEPS=1 bash scripts/train.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_AMP_8GPU.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_AMP_8GPU.sh
new file mode 100644
index 0000000000..8953566d4a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_AMP_8GPU.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+NUM_GPUS=8 AMP=true BATCH_SIZE=64 GRADIENT_ACCUMULATION_STEPS=1 bash scripts/train.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_FP32_16GPU.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_FP32_16GPU.sh
new file mode 100644
index 0000000000..873fb92f17
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_FP32_16GPU.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+NUM_GPUS=16 BATCH_SIZE=64 GRADIENT_ACCUMULATION_STEPS=1 bash scripts/train.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_FP32_8GPU.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_FP32_8GPU.sh
new file mode 100644
index 0000000000..4ac61b677d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGX2_Jasper_FP32_8GPU.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+NUM_GPUS=8 AMP=true BATCH_SIZE=64 GRADIENT_ACCUMULATION_STEPS=2 bash scripts/train.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGXA100_Jasper_AMP_8GPU.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGXA100_Jasper_AMP_8GPU.sh
new file mode 100644
index 0000000000..8953566d4a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGXA100_Jasper_AMP_8GPU.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+NUM_GPUS=8 AMP=true BATCH_SIZE=64 GRADIENT_ACCUMULATION_STEPS=1 bash scripts/train.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGXA100_Jasper_TF32_8GPU.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGXA100_Jasper_TF32_8GPU.sh
new file mode 100644
index 0000000000..eed6a1273f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/platform/DGXA100_Jasper_TF32_8GPU.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+NUM_GPUS=8 BATCH_SIZE=64 GRADIENT_ACCUMULATION_STEPS=2 bash scripts/train.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/requirements.txt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/requirements.txt
new file mode 100644
index 0000000000..92eb868b1c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/requirements.txt
@@ -0,0 +1,13 @@
+ascii-graph==1.5.1
+inflect==5.3.0
+ipdb
+librosa==0.8.0
+pandas==1.1.4
+pycuda==2020.1
+pyyaml>=5.4
+soundfile
+sox==1.4.1
+tqdm==4.53.0
+unidecode==1.2.0
+wrapt==1.10.11
+git+git://github.com/NVIDIA/dllogger.git@26a0f8f1958de2c0c460925ff6102a4d2486d6cc#egg=dllogger
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/docker/build.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/docker/build.sh
new file mode 100644
index 0000000000..cfdc97c010
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/docker/build.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+docker build . --rm -t jasper
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/docker/launch.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/docker/launch.sh
new file mode 100644
index 0000000000..bc719ffadc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/docker/launch.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+
+SCRIPT_DIR=$(cd $(dirname $0); pwd)
+: ${JASPER_REPO:="$SCRIPT_DIR/../.."}
+
+: ${DATA_DIR:=${1:-"$JASPER_REPO/datasets"}}
+: ${CHECKPOINT_DIR:=${2:-"$JASPER_REPO/checkpoints"}}
+: ${OUTPUT_DIR:=${3:-"$JASPER_REPO/results"}}
+: ${SCRIPT:=${4:-}}
+
+mkdir -p $DATA_DIR
+mkdir -p $CHECKPOINT_DIR
+mkdir -p $OUTPUT_DIR
+
+MOUNTS=""
+MOUNTS+=" -v $DATA_DIR:/dataset"
+MOUNTS+=" -v $CHECKPOINT_DIR:/checkpoints"
+MOUNTS+=" -v $OUTPUT_DIR:/results"
+MOUNTS+=" -v $JASPER_REPO:/workspace/jasper"
+MOUNTS+=" -v /usr/local/Ascend:/usr/local/Ascend"
+
+echo $MOUNTS
+docker run -it --rm --device /dev/davinci5 --device /dev/davinci_manager --device /dev/hisi_hdc --device /dev/devmm_svm \
+  --env PYTHONDONTWRITEBYTECODE=1 \
+  --shm-size=4g \
+  --ulimit memlock=-1 \
+  --ulimit stack=67108864 \
+  $MOUNTS \
+  $EXTRA_JASPER_ENV \
+  -w /workspace/jasper \
+  jasper:latest bash $SCRIPT
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/download_librispeech.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/download_librispeech.sh
new file mode 100644
index 0000000000..b96744eb94
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/download_librispeech.sh
@@ -0,0 +1,32 @@
+#!/usr/bin/env bash
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+
+DATA_SET="LibriSpeech"
+DATA_ROOT_DIR="/home/cyl_dataset"
+DATA_DIR="${DATA_ROOT_DIR}/${DATA_SET}"
+
+if [ ! -d "$DATA_DIR" ]
+then
+   mkdir --mode 755 $DATA_DIR
+
+   python utils/download_librispeech.py \
+      utils/librispeech.csv \
+      $DATA_DIR \
+      -e ${DATA_ROOT_DIR}/
+else
+   echo "Directory $DATA_DIR already exists."
+fi
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/evaluation.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/evaluation.sh
new file mode 100644
index 0000000000..08009e514e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/evaluation.sh
@@ -0,0 +1,22 @@
+#!/bin/bash
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+set -a
+
+: ${PREDICTION_FILE:=}
+: ${DATASET:="test-other"}
+
+bash ./scripts/inference.sh "$@"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/inference.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/inference.sh
new file mode 100644
index 0000000000..f60c3d5875
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/inference.sh
@@ -0,0 +1,65 @@
+#!/bin/bash
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+: ${DATA_DIR:=${1:-"/datasets/LibriSpeech"}}
+: ${MODEL_CONFIG:=${2:-"configs/jasper10x5dr_speedp-online_speca.yaml"}}
+: ${OUTPUT_DIR:=${3:-"/results"}}
+: ${CHECKPOINT:=${4:-"/checkpoints/jasper_fp16.pt"}}
+: ${DATASET:="test-other"}
+: ${LOG_FILE:=""}
+: ${CUDNN_BENCHMARK:=false}
+: ${MAX_DURATION:=""}
+: ${PAD_TO_MAX_DURATION:=false}
+: ${PAD_LEADING:=16}
+: ${NUM_GPUS:=1}
+: ${NUM_STEPS:=0}
+: ${NUM_WARMUP_STEPS:=0}
+: ${AMP:=false}
+: ${BATCH_SIZE:=64}
+: ${EMA:=true}
+: ${SEED:=0}
+: ${DALI_DEVICE:="gpu"}
+: ${CPU:=false}
+: ${LOGITS_FILE:=}
+: ${PREDICTION_FILE:="${OUTPUT_DIR}/${DATASET}.predictions"}
+
+mkdir -p "$OUTPUT_DIR"
+
+ARGS="--dataset_dir=$DATA_DIR"
+ARGS+=" --val_manifest=$DATA_DIR/librispeech-${DATASET}-wav.json"
+ARGS+=" --model_config=$MODEL_CONFIG"
+ARGS+=" --output_dir=$OUTPUT_DIR"
+ARGS+=" --batch_size=$BATCH_SIZE"
+ARGS+=" --seed=$SEED"
+ARGS+=" --dali_device=$DALI_DEVICE"
+ARGS+=" --steps $NUM_STEPS"
+ARGS+=" --warmup_steps $NUM_WARMUP_STEPS"
+ARGS+=" --pad_leading $PAD_LEADING"
+
+[ "$AMP" = true ] &&                 ARGS+=" --amp"
+[ "$EMA" = true ] &&                 ARGS+=" --ema"
+[ "$CUDNN_BENCHMARK" = true ] &&     ARGS+=" --cudnn_benchmark"
+[ -n "$CHECKPOINT" ] &&              ARGS+=" --ckpt=${CHECKPOINT}"
+[ -n "$LOG_FILE" ] &&                ARGS+=" --log_file $LOG_FILE"
+[ -n "$PREDICTION_FILE" ] &&         ARGS+=" --save_prediction $PREDICTION_FILE"
+[ -n "$LOGITS_FILE" ] &&             ARGS+=" --logits_save_to $LOGITS_FILE"
+[ "$CPU" == "true" ] &&              ARGS+=" --cpu"
+[ -n "$MAX_DURATION" ] &&            ARGS+=" --override_config input_val.audio_dataset.max_duration=$MAX_DURATION" \
+                                     ARGS+=" --override_config input_val.filterbank_features.max_duration=$MAX_DURATION"
+[ "$PAD_TO_MAX_DURATION" = true ] && ARGS+=" --override_config input_val.audio_dataset.pad_to_max_duration=True" \
+                                     ARGS+=" --override_config input_val.filterbank_features.pad_to_max_duration=True"
+
+python -m torch.distributed.launch --nproc_per_node=$NUM_GPUS inference.py $ARGS
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/inference_benchmark.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/inference_benchmark.sh
new file mode 100644
index 0000000000..66dfe5ab72
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/inference_benchmark.sh
@@ -0,0 +1,38 @@
+#!/bin/bash
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+set -a
+
+: ${OUTPUT_DIR:=${3:-"/results"}}
+: ${CUDNN_BENCHMARK:=true}
+: ${PAD_TO_MAX_DURATION:=true}
+: ${PAD_LEADING:=0}
+: ${NUM_WARMUP_STEPS:=10}
+: ${NUM_STEPS:=500}
+
+: ${AMP:=false}
+: ${DALI_DEVICE:="cpu"}
+: ${BATCH_SIZE_SEQ:="1 2 4 8 16"}
+: ${MAX_DURATION_SEQ:="2 7 16.7"}
+
+for MAX_DURATION in $MAX_DURATION_SEQ; do
+  for BATCH_SIZE in $BATCH_SIZE_SEQ; do
+
+    LOG_FILE="$OUTPUT_DIR/perf-infer_dali-${DALI_DEVICE}_amp-${AMP}_dur${MAX_DURATION}_bs${BATCH_SIZE}.json"
+    bash ./scripts/inference.sh "$@"
+
+  done
+done
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/preprocess_librispeech.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/preprocess_librispeech.sh
new file mode 100644
index 0000000000..3ae5b5768f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/preprocess_librispeech.sh
@@ -0,0 +1,54 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+#!/usr/bin/env bash
+
+SPEEDS=$1
+[ -n "$SPEEDS" ] && SPEED_FLAG="--speed $SPEEDS"
+
+python ./utils/convert_librispeech.py \
+    --input_dir /home/dataset/LibriSpeech/train-clean-100 \
+    --dest_dir /home/dataset/LibriSpeech/train-clean-100-wav \
+    --output_json /home/dataset/LibriSpeech/librispeech-train-clean-100-wav.json \
+    $SPEED_FLAG
+python ./utils/convert_librispeech.py \
+    --input_dir /home/dataset/LibriSpeech/train-clean-360 \
+    --dest_dir /home/dataset/LibriSpeech/train-clean-360-wav \
+    --output_json /home/dataset/LibriSpeech/librispeech-train-clean-360-wav.json \
+    $SPEED_FLAG
+python ./utils/convert_librispeech.py \
+    --input_dir /home/dataset/LibriSpeech/train-other-500 \
+    --dest_dir /home/dataset/LibriSpeech/train-other-500-wav \
+    --output_json /home/dataset/LibriSpeech/librispeech-train-other-500-wav.json \
+    $SPEED_FLAG
+
+
+python ./utils/convert_librispeech.py \
+    --input_dir /home/dataset/LibriSpeech/dev-clean \
+    --dest_dir /home/dataset/LibriSpeech/dev-clean-wav \
+    --output_json /home/dataset/LibriSpeech/librispeech-dev-clean-wav.json
+python ./utils/convert_librispeech.py \
+    --input_dir /home/dataset/LibriSpeech/dev-other \
+    --dest_dir /home/dataset/LibriSpeech/dev-other-wav \
+    --output_json /home/dataset/LibriSpeech/librispeech-dev-other-wav.json
+
+
+python ./utils/convert_librispeech.py \
+    --input_dir /home/dataset/LibriSpeech/test-clean \
+    --dest_dir /home/dataset/LibriSpeech/test-clean-wav \
+    --output_json /home/dataset/LibriSpeech/librispeech-test-clean-wav.json
+python ./utils/convert_librispeech.py \
+    --input_dir /home/dataset/LibriSpeech/test-other \
+    --dest_dir /home/dataset/LibriSpeech/test-other-wav \
+    --output_json /home/dataset/LibriSpeech/librispeech-test-other-wav.json
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/train.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/train.sh
new file mode 100644
index 0000000000..51ad11e924
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/train.sh
@@ -0,0 +1,90 @@
+#!/bin/bash
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+export OMP_NUM_THREADS=1
+
+: ${DATA_DIR:=${1:-"/datasets/LibriSpeech"}}
+: ${MODEL_CONFIG:=${2:-"configs/jasper10x5dr_speedp-online_speca.yaml"}}
+: ${OUTPUT_DIR:=${3:-"/results"}}
+: ${CHECKPOINT:=${4:-}}
+: ${RESUME:=true}
+: ${CUDNN_BENCHMARK:=true}
+: ${NUM_GPUS:=8}
+: ${AMP:=True}
+: ${BATCH_SIZE:=32}
+: ${GRAD_ACCUMULATION_STEPS:=2}
+: ${LEARNING_RATE:=0.01}
+: ${MIN_LEARNING_RATE:=0.00001}
+: ${LR_POLICY:=exponential}
+: ${LR_EXP_GAMMA:=0.981}
+: ${EMA:=0.999}
+: ${SEED:=0}
+: ${EPOCHS:=440}
+: ${WARMUP_EPOCHS:=2}
+: ${HOLD_EPOCHS:=140}
+: ${SAVE_FREQUENCY:=10}
+: ${EPOCHS_THIS_JOB:=0}
+: ${DALI_DEVICE:="cpu"}
+: ${PAD_TO_MAX_DURATION:=false}
+: ${EVAL_FREQUENCY:=544}
+: ${PREDICTION_FREQUENCY:=544}
+: ${TRAIN_MANIFESTS:="$DATA_DIR/librispeech-train-clean-100-wav.json \
+                      $DATA_DIR/librispeech-train-clean-360-wav.json \
+                      $DATA_DIR/librispeech-train-other-500-wav.json"}
+: ${VAL_MANIFESTS:="$DATA_DIR/librispeech-dev-clean-wav.json"}
+
+mkdir -p "$OUTPUT_DIR"
+
+ARGS="--dataset_dir=$DATA_DIR"
+ARGS+=" --val_manifests $VAL_MANIFESTS"
+ARGS+=" --train_manifests $TRAIN_MANIFESTS"
+ARGS+=" --model_config=$MODEL_CONFIG"
+ARGS+=" --output_dir=$OUTPUT_DIR"
+ARGS+=" --lr=$LEARNING_RATE"
+ARGS+=" --batch_size=$BATCH_SIZE"
+ARGS+=" --min_lr=$MIN_LEARNING_RATE"
+ARGS+=" --lr_policy=$LR_POLICY"
+ARGS+=" --lr_exp_gamma=$LR_EXP_GAMMA"
+ARGS+=" --epochs=$EPOCHS"
+ARGS+=" --warmup_epochs=$WARMUP_EPOCHS"
+ARGS+=" --hold_epochs=$HOLD_EPOCHS"
+ARGS+=" --epochs_this_job=$EPOCHS_THIS_JOB"
+ARGS+=" --ema=$EMA"
+ARGS+=" --seed=$SEED"
+ARGS+=" --optimizer=novograd"
+ARGS+=" --weight_decay=1e-3"
+ARGS+=" --save_frequency=$SAVE_FREQUENCY"
+ARGS+=" --keep_milestones 100 200 300 400"
+ARGS+=" --save_best_from=380"
+ARGS+=" --log_frequency=1"
+ARGS+=" --eval_frequency=$EVAL_FREQUENCY"
+ARGS+=" --prediction_frequency=$PREDICTION_FREQUENCY"
+ARGS+=" --grad_accumulation_steps=$GRAD_ACCUMULATION_STEPS "
+ARGS+=" --dali_device=$DALI_DEVICE"
+
+[ "$AMP" = true ] &&                 ARGS+=" --amp"
+[ "$RESUME" = true ] &&              ARGS+=" --resume"
+[ "$CUDNN_BENCHMARK" = true ] &&     ARGS+=" --cudnn_benchmark"
+[ -n "$MAX_DURATION" ] &&            ARGS+=" --override_config input_train.audio_dataset.max_duration=$MAX_DURATION" \
+                                     ARGS+=" --override_config input_train.filterbank_features.max_duration=$MAX_DURATION"
+[ "$PAD_TO_MAX_DURATION" = true ] && ARGS+=" --override_config input_train.audio_dataset.pad_to_max_duration=True" \
+                                     ARGS+=" --override_config input_train.filterbank_features.pad_to_max_duration=True"
+[ -n "$CHECKPOINT" ] &&              ARGS+=" --ckpt=$CHECKPOINT"
+[ -n "$LOG_FILE" ] &&                ARGS+=" --log_file $LOG_FILE"
+[ -n "$PRE_ALLOCATE" ] &&            ARGS+=" --pre_allocate_range $PRE_ALLOCATE"
+
+DISTRIBUTED="-m torch.distributed.launch --nproc_per_node=$NUM_GPUS"
+python $DISTRIBUTED train.py $ARGS
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/train_benchmark.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/train_benchmark.sh
new file mode 100644
index 0000000000..f70760fe58
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/scripts/train_benchmark.sh
@@ -0,0 +1,49 @@
+#!/bin/bash
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+set -a
+
+# measure on speed perturbed data, but so slightly that fbank length remains the same
+# with pad_to_max_duration, this reduces cuDNN benchmak's burn-in period to a single step
+: ${DATA_DIR:=${1:-"/datasets/LibriSpeech"}}
+: ${OUTPUT_DIR:=${3:-"/results"}}
+: ${TRAIN_MANIFESTS:="$DATA_DIR/librispeech-train-clean-100-wav.json"}
+
+# run for a number of epochs, but don't finalize the training
+: ${EPOCHS_THIS_JOB:=2}
+: ${EPOCHS:=100000}
+: ${RESUME:=false}
+: ${SAVE_FREQUENCY:=100000}
+: ${EVAL_FREQUENCY:=100000}
+: ${GRAD_ACCUMULATION_STEPS:=1}
+
+: ${AMP:=false}
+: ${EMA:=0}
+: ${DALI_DEVICE:="gpu"}
+: ${NUM_GPUS_SEQ:="1 4 8"}
+: ${BATCH_SIZE_SEQ:="32"}
+# A probable range of batch lengths for LibriSpeech
+# with BS=64 and continuous speed perturbation (0.85, 1.15)
+: ${PRE_ALLOCATE:="1408 1920"}
+
+for NUM_GPUS in $NUM_GPUS_SEQ; do
+  for BATCH_SIZE in $BATCH_SIZE_SEQ; do
+
+    LOG_FILE="$OUTPUT_DIR/perf-train_dali-${DALI_DEVICE}_amp-${AMP}_ngpus${NUM_GPUS}_bs${BATCH_SIZE}.json"
+    bash ./scripts/train.sh "$@"
+
+  done
+done
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/docker/build.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/docker/build.sh
new file mode 100644
index 0000000000..cfdc97c010
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/docker/build.sh
@@ -0,0 +1,3 @@
+#!/bin/bash
+
+docker build . --rm -t jasper
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/docker/launch.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/docker/launch.sh
new file mode 100644
index 0000000000..bc719ffadc
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/docker/launch.sh
@@ -0,0 +1,31 @@
+#!/bin/bash
+
+SCRIPT_DIR=$(cd $(dirname $0); pwd)
+: ${JASPER_REPO:="$SCRIPT_DIR/../.."}
+
+: ${DATA_DIR:=${1:-"$JASPER_REPO/datasets"}}
+: ${CHECKPOINT_DIR:=${2:-"$JASPER_REPO/checkpoints"}}
+: ${OUTPUT_DIR:=${3:-"$JASPER_REPO/results"}}
+: ${SCRIPT:=${4:-}}
+
+mkdir -p $DATA_DIR
+mkdir -p $CHECKPOINT_DIR
+mkdir -p $OUTPUT_DIR
+
+MOUNTS=""
+MOUNTS+=" -v $DATA_DIR:/dataset"
+MOUNTS+=" -v $CHECKPOINT_DIR:/checkpoints"
+MOUNTS+=" -v $OUTPUT_DIR:/results"
+MOUNTS+=" -v $JASPER_REPO:/workspace/jasper"
+MOUNTS+=" -v /usr/local/Ascend:/usr/local/Ascend"
+
+echo $MOUNTS
+docker run -it --rm --device /dev/davinci5 --device /dev/davinci_manager --device /dev/hisi_hdc --device /dev/devmm_svm \
+  --env PYTHONDONTWRITEBYTECODE=1 \
+  --shm-size=4g \
+  --ulimit memlock=-1 \
+  --ulimit stack=67108864 \
+  $MOUNTS \
+  $EXTRA_JASPER_ENV \
+  -w /workspace/jasper \
+  jasper:latest bash $SCRIPT
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/env_npu.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/env_npu.sh
new file mode 100644
index 0000000000..baf2d93543
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/env_npu.sh
@@ -0,0 +1,80 @@
+#!/bin/bash
+export install_path=/usr/local/Ascend
+
+if [ -d ${install_path}/toolkit ]; then
+    export LD_LIBRARY_PATH=${install_path}/fwkacllib/lib64/:/usr/include/hdf5/lib/:/usr/local/:/usr/local/lib/:/usr/lib/:${install_path}/driver/lib64/common/:${install_path}/driver/lib64/driver/:${install_path}/add-ons:${path_lib}:${LD_LIBRARY_PATH}
+    export PATH=${install_path}/fwkacllib/ccec_compiler/bin:${install_path}/fwkacllib/bin:$PATH
+    export PYTHONPATH=${install_path}/fwkacllib/python/site-packages:${install_path}/tfplugin/python/site-packages:${install_path}/toolkit/python/site-packages:$PYTHONPATH
+    export PYTHONPATH=/usr/local/python3.7.5/lib/python3.7/site-packages:$PYTHONPATH
+    export ASCEND_OPP_PATH=${install_path}/opp
+else
+    if [ -d ${install_path}/nnae/latest ];then
+        export LD_LIBRARY_PATH=${install_path}/nnae/latest/fwkacllib/lib64/:/usr/local/:/usr/local/python3.7.5/lib/:/usr/local/openblas/lib:/usr/local/lib/:/usr/lib64/:/usr/lib/:${install_path}/driver/lib64/common/:${install_path}/driver/lib64/driver/:${install_path}/add-ons/:/usr/lib/aarch64_64-linux-gnu:$LD_LIBRARY_PATH
+        export PATH=$PATH:${install_path}/nnae/latest/fwkacllib/ccec_compiler/bin/:${install_path}/nnae/latest/toolkit/tools/ide_daemon/bin/
+        export ASCEND_OPP_PATH=${install_path}/nnae/latest/opp/
+        export OPTION_EXEC_EXTERN_PLUGIN_PATH=${install_path}/nnae/latest/fwkacllib/lib64/plugin/opskernel/libfe.so:${install_path}/nnae/latest/fwkacllib/lib64/plugin/opskernel/libaicpu_engine.so:${install_path}/nnae/latest/fwkacllib/lib64/plugin/opskernel/libge_local_engine.so
+        export PYTHONPATH=${install_path}/nnae/latest/fwkacllib/python/site-packages/:${install_path}/nnae/latest/fwkacllib/python/site-packages/auto_tune.egg/auto_tune:${install_path}/nnae/latest/fwkacllib/python/site-packages/schedule_search.egg:$PYTHONPATH
+        export ASCEND_AICPU_PATH=${install_path}/nnae/latest
+    else
+        export LD_LIBRARY_PATH=${install_path}/ascend-toolkit/latest/fwkacllib/lib64/:/usr/local/:/usr/local/lib/:/usr/lib64/:/usr/lib/:/usr/local/python3.7.5/lib/:/usr/local/openblas/lib:${install_path}/driver/lib64/common/:${install_path}/driver/lib64/driver/:${install_path}/add-ons/:/usr/lib/aarch64-linux-gnu:$LD_LIBRARY_PATH
+        export PATH=$PATH:${install_path}/ascend-toolkit/latest/fwkacllib/ccec_compiler/bin/:${install_path}/ascend-toolkit/latest/toolkit/tools/ide_daemon/bin/
+        export ASCEND_OPP_PATH=${install_path}/ascend-toolkit/latest/opp/
+        export OPTION_EXEC_EXTERN_PLUGIN_PATH=${install_path}/ascend-toolkit/latest/fwkacllib/lib64/plugin/opskernel/libfe.so:${install_path}/ascend-toolkit/latest/fwkacllib/lib64/plugin/opskernel/libaicpu_engine.so:${install_path}/ascend-toolkit/latest/fwkacllib/lib64/plugin/opskernel/libge_local_engine.so
+        export PYTHONPATH=${install_path}/ascend-toolkit/latest/fwkacllib/python/site-packages/:${install_path}/ascend-toolkit/latest/fwkacllib/python/site-packages/auto_tune.egg/auto_tune:${install_path}/ascend-toolkit/latest/fwkacllib/python/site-packages/schedule_search.egg:$PYTHONPATH
+        export ASCEND_AICPU_PATH=${install_path}/ascend-toolkit/latest
+    fi
+fi
+
+${install_path}/driver/tools/msnpureport -g error -d 0
+${install_path}/driver/tools/msnpureport -g error -d 1
+${install_path}/driver/tools/msnpureport -g error -d 2
+${install_path}/driver/tools/msnpureport -g error -d 3
+${install_path}/driver/tools/msnpureport -g error -d 4
+${install_path}/driver/tools/msnpureport -g error -d 5
+${install_path}/driver/tools/msnpureport -g error -d 6
+${install_path}/driver/tools/msnpureport -g error -d 7
+
+#½«HostÈÕÖ¾Êä³öµ½´®¿Ú,0-¹Ø±Õ/1-¿ªÆô
+export ASCEND_SLOG_PRINT_TO_STDOUT=0
+#ÉèÖÃĬÈÏÈÕÖ¾¼¶±ð,0-debug/1-info/2-warning/3-error
+export ASCEND_GLOBAL_LOG_LEVEL=3
+#ÉèÖÃEventÈÕÖ¾¿ªÆô±êÖ¾,0-¹Ø±Õ/1-¿ªÆô
+export ASCEND_GLOBAL_EVENT_ENABLE=0
+#ÉèÖÃÊÇ·ñ¿ªÆôtaskque,0-¹Ø±Õ/1-¿ªÆô
+export TASK_QUEUE_ENABLE=1
+#ÉèÖÃÊÇ·ñ¿ªÆôPTCopy,0-¹Ø±Õ/1-¿ªÆô
+export PTCOPY_ENABLE=1
+#ÉèÖÃÊÇ·ñ¿ªÆô2¸ö·ÇÁ¬Ðøcombined±êÖ¾,0-¹Ø±Õ/1-¿ªÆô
+export COMBINED_ENABLE=1
+#ÉèÖÃÊÇ·ñ¿ªÆô3¸ö·ÇÁ¬Ðøcombined±êÖ¾,0-¹Ø±Õ/1-¿ªÆô
+export TRI_COMBINED_ENABLE=1
+#ÉèÖÃÌØÊⳡ¾°ÊÇ·ñÐèÒªÖØбàÒë,²»ÐèÒªÐÞ¸Ä
+export DYNAMIC_OP="ADD#MUL"
+# HCCL°×Ãûµ¥¿ª¹Ø,1-¹Ø±Õ/0-¿ªÆô
+export HCCL_WHITELIST_DISABLE=1
+# HCCLĬÈϳ¬Ê±Ê±¼ä120s½ÏÉÙ£¬ÐÞ¸ÄΪ1800s¶ÔÆëPyTorchĬÈÏÉèÖÃ
+export HCCL_CONNECT_TIMEOUT=1800
+
+ulimit -SHn 512000
+
+path_lib=$(python3.7 -c """
+import sys
+import re
+result=''
+for index in range(len(sys.path)):
+    match_sit = re.search('-packages', sys.path[index])
+    if match_sit is not None:
+        match_lib = re.search('lib', sys.path[index])
+
+        if match_lib is not None:
+            end=match_lib.span()[1]
+            result += sys.path[index][0:end] + ':'
+
+        result+=sys.path[index] + '/torch/lib:'
+print(result)"""
+)
+
+echo ${path_lib}
+
+export LD_LIBRARY_PATH=/usr/local/python3.7.5/lib/:${path_lib}:$LD_LIBRARY_PATH
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_eval_8p.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_eval_8p.sh
new file mode 100644
index 0000000000..f60c3d5875
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_eval_8p.sh
@@ -0,0 +1,65 @@
+#!/bin/bash
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+: ${DATA_DIR:=${1:-"/datasets/LibriSpeech"}}
+: ${MODEL_CONFIG:=${2:-"configs/jasper10x5dr_speedp-online_speca.yaml"}}
+: ${OUTPUT_DIR:=${3:-"/results"}}
+: ${CHECKPOINT:=${4:-"/checkpoints/jasper_fp16.pt"}}
+: ${DATASET:="test-other"}
+: ${LOG_FILE:=""}
+: ${CUDNN_BENCHMARK:=false}
+: ${MAX_DURATION:=""}
+: ${PAD_TO_MAX_DURATION:=false}
+: ${PAD_LEADING:=16}
+: ${NUM_GPUS:=1}
+: ${NUM_STEPS:=0}
+: ${NUM_WARMUP_STEPS:=0}
+: ${AMP:=false}
+: ${BATCH_SIZE:=64}
+: ${EMA:=true}
+: ${SEED:=0}
+: ${DALI_DEVICE:="gpu"}
+: ${CPU:=false}
+: ${LOGITS_FILE:=}
+: ${PREDICTION_FILE:="${OUTPUT_DIR}/${DATASET}.predictions"}
+
+mkdir -p "$OUTPUT_DIR"
+
+ARGS="--dataset_dir=$DATA_DIR"
+ARGS+=" --val_manifest=$DATA_DIR/librispeech-${DATASET}-wav.json"
+ARGS+=" --model_config=$MODEL_CONFIG"
+ARGS+=" --output_dir=$OUTPUT_DIR"
+ARGS+=" --batch_size=$BATCH_SIZE"
+ARGS+=" --seed=$SEED"
+ARGS+=" --dali_device=$DALI_DEVICE"
+ARGS+=" --steps $NUM_STEPS"
+ARGS+=" --warmup_steps $NUM_WARMUP_STEPS"
+ARGS+=" --pad_leading $PAD_LEADING"
+
+[ "$AMP" = true ] &&                 ARGS+=" --amp"
+[ "$EMA" = true ] &&                 ARGS+=" --ema"
+[ "$CUDNN_BENCHMARK" = true ] &&     ARGS+=" --cudnn_benchmark"
+[ -n "$CHECKPOINT" ] &&              ARGS+=" --ckpt=${CHECKPOINT}"
+[ -n "$LOG_FILE" ] &&                ARGS+=" --log_file $LOG_FILE"
+[ -n "$PREDICTION_FILE" ] &&         ARGS+=" --save_prediction $PREDICTION_FILE"
+[ -n "$LOGITS_FILE" ] &&             ARGS+=" --logits_save_to $LOGITS_FILE"
+[ "$CPU" == "true" ] &&              ARGS+=" --cpu"
+[ -n "$MAX_DURATION" ] &&            ARGS+=" --override_config input_val.audio_dataset.max_duration=$MAX_DURATION" \
+                                     ARGS+=" --override_config input_val.filterbank_features.max_duration=$MAX_DURATION"
+[ "$PAD_TO_MAX_DURATION" = true ] && ARGS+=" --override_config input_val.audio_dataset.pad_to_max_duration=True" \
+                                     ARGS+=" --override_config input_val.filterbank_features.pad_to_max_duration=True"
+
+python -m torch.distributed.launch --nproc_per_node=$NUM_GPUS inference.py $ARGS
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_full_1p.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_full_1p.sh
new file mode 100644
index 0000000000..efd798f999
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_full_1p.sh
@@ -0,0 +1,91 @@
+#!/bin/bash
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+export OMP_NUM_THREADS=1
+
+: ${DATA_DIR:=${1:-"/datasets/LibriSpeech"}}
+: ${MODEL_CONFIG:=${2:-"configs/jasper10x5dr_speedp-online_speca.yaml"}}
+: ${OUTPUT_DIR:=${3:-"/results"}}
+: ${CHECKPOINT:=${4:-}}
+: ${RESUME:=true}
+: ${CUDNN_BENCHMARK:=true}
+: ${NUM_GPUS:=1}
+: ${AMP:=True}
+: ${BATCH_SIZE:=32}
+: ${GRAD_ACCUMULATION_STEPS:=2}
+: ${LEARNING_RATE:=0.01}
+: ${MIN_LEARNING_RATE:=0.00001}
+: ${LR_POLICY:=exponential}
+: ${LR_EXP_GAMMA:=0.981}
+: ${EMA:=0.999}
+: ${SEED:=0}
+: ${EPOCHS:=440}
+: ${WARMUP_EPOCHS:=2}
+: ${HOLD_EPOCHS:=140}
+: ${SAVE_FREQUENCY:=10}
+: ${EPOCHS_THIS_JOB:=0}
+: ${DALI_DEVICE:="cpu"}
+: ${PAD_TO_MAX_DURATION:=false}
+: ${EVAL_FREQUENCY:=544}
+: ${PREDICTION_FREQUENCY:=544}
+: ${TRAIN_MANIFESTS:="$DATA_DIR/librispeech-train-clean-100-wav.json \
+                      $DATA_DIR/librispeech-train-clean-360-wav.json \
+                      $DATA_DIR/librispeech-train-other-500-wav.json"}
+: ${VAL_MANIFESTS:="$DATA_DIR/librispeech-dev-clean-wav.json"}
+
+mkdir -p "$OUTPUT_DIR"
+
+ARGS="--dataset_dir=$DATA_DIR"
+ARGS+=" --val_manifests $VAL_MANIFESTS"
+ARGS+=" --train_manifests $TRAIN_MANIFESTS"
+ARGS+=" --model_config=$MODEL_CONFIG"
+ARGS+=" --output_dir=$OUTPUT_DIR"
+ARGS+=" --lr=$LEARNING_RATE"
+ARGS+=" --batch_size=$BATCH_SIZE"
+ARGS+=" --min_lr=$MIN_LEARNING_RATE"
+ARGS+=" --lr_policy=$LR_POLICY"
+ARGS+=" --lr_exp_gamma=$LR_EXP_GAMMA"
+ARGS+=" --epochs=$EPOCHS"
+ARGS+=" --warmup_epochs=$WARMUP_EPOCHS"
+ARGS+=" --hold_epochs=$HOLD_EPOCHS"
+ARGS+=" --epochs_this_job=$EPOCHS_THIS_JOB"
+ARGS+=" --ema=$EMA"
+ARGS+=" --seed=$SEED"
+ARGS+=" --optimizer=novograd"
+ARGS+=" --weight_decay=1e-3"
+ARGS+=" --save_frequency=$SAVE_FREQUENCY"
+ARGS+=" --keep_milestones 100 200 300 400"
+ARGS+=" --save_best_from=380"
+ARGS+=" --log_frequency=1"
+ARGS+=" --eval_frequency=$EVAL_FREQUENCY"
+ARGS+=" --prediction_frequency=$PREDICTION_FREQUENCY"
+ARGS+=" --grad_accumulation_steps=$GRAD_ACCUMULATION_STEPS "
+ARGS+=" --dali_device=$DALI_DEVICE"
+
+[ "$AMP" = true ] &&                 ARGS+=" --amp"
+[ "$RESUME" = true ] &&              ARGS+=" --resume"
+[ "$CUDNN_BENCHMARK" = true ] &&     ARGS+=" --cudnn_benchmark"
+[ -n "$MAX_DURATION" ] &&            ARGS+=" --override_config input_train.audio_dataset.max_duration=$MAX_DURATION" \
+                                     ARGS+=" --override_config input_train.filterbank_features.max_duration=$MAX_DURATION"
+[ "$PAD_TO_MAX_DURATION" = true ] && ARGS+=" --override_config input_train.audio_dataset.pad_to_max_duration=True" \
+                                     ARGS+=" --override_config input_train.filterbank_features.pad_to_max_duration=True"
+[ -n "$CHECKPOINT" ] &&              ARGS+=" --ckpt=$CHECKPOINT"
+[ -n "$LOG_FILE" ] &&                ARGS+=" --log_file $LOG_FILE"
+[ -n "$PRE_ALLOCATE" ] &&            ARGS+=" --pre_allocate_range $PRE_ALLOCATE"
+
+DISTRIBUTED="-m torch.distributed.launch --nproc_per_node=$NUM_GPUS"
+python $DISTRIBUTED train.py $ARGS
+
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_full_8p.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_full_8p.sh
new file mode 100644
index 0000000000..51ad11e924
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_full_8p.sh
@@ -0,0 +1,90 @@
+#!/bin/bash
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+export OMP_NUM_THREADS=1
+
+: ${DATA_DIR:=${1:-"/datasets/LibriSpeech"}}
+: ${MODEL_CONFIG:=${2:-"configs/jasper10x5dr_speedp-online_speca.yaml"}}
+: ${OUTPUT_DIR:=${3:-"/results"}}
+: ${CHECKPOINT:=${4:-}}
+: ${RESUME:=true}
+: ${CUDNN_BENCHMARK:=true}
+: ${NUM_GPUS:=8}
+: ${AMP:=True}
+: ${BATCH_SIZE:=32}
+: ${GRAD_ACCUMULATION_STEPS:=2}
+: ${LEARNING_RATE:=0.01}
+: ${MIN_LEARNING_RATE:=0.00001}
+: ${LR_POLICY:=exponential}
+: ${LR_EXP_GAMMA:=0.981}
+: ${EMA:=0.999}
+: ${SEED:=0}
+: ${EPOCHS:=440}
+: ${WARMUP_EPOCHS:=2}
+: ${HOLD_EPOCHS:=140}
+: ${SAVE_FREQUENCY:=10}
+: ${EPOCHS_THIS_JOB:=0}
+: ${DALI_DEVICE:="cpu"}
+: ${PAD_TO_MAX_DURATION:=false}
+: ${EVAL_FREQUENCY:=544}
+: ${PREDICTION_FREQUENCY:=544}
+: ${TRAIN_MANIFESTS:="$DATA_DIR/librispeech-train-clean-100-wav.json \
+                      $DATA_DIR/librispeech-train-clean-360-wav.json \
+                      $DATA_DIR/librispeech-train-other-500-wav.json"}
+: ${VAL_MANIFESTS:="$DATA_DIR/librispeech-dev-clean-wav.json"}
+
+mkdir -p "$OUTPUT_DIR"
+
+ARGS="--dataset_dir=$DATA_DIR"
+ARGS+=" --val_manifests $VAL_MANIFESTS"
+ARGS+=" --train_manifests $TRAIN_MANIFESTS"
+ARGS+=" --model_config=$MODEL_CONFIG"
+ARGS+=" --output_dir=$OUTPUT_DIR"
+ARGS+=" --lr=$LEARNING_RATE"
+ARGS+=" --batch_size=$BATCH_SIZE"
+ARGS+=" --min_lr=$MIN_LEARNING_RATE"
+ARGS+=" --lr_policy=$LR_POLICY"
+ARGS+=" --lr_exp_gamma=$LR_EXP_GAMMA"
+ARGS+=" --epochs=$EPOCHS"
+ARGS+=" --warmup_epochs=$WARMUP_EPOCHS"
+ARGS+=" --hold_epochs=$HOLD_EPOCHS"
+ARGS+=" --epochs_this_job=$EPOCHS_THIS_JOB"
+ARGS+=" --ema=$EMA"
+ARGS+=" --seed=$SEED"
+ARGS+=" --optimizer=novograd"
+ARGS+=" --weight_decay=1e-3"
+ARGS+=" --save_frequency=$SAVE_FREQUENCY"
+ARGS+=" --keep_milestones 100 200 300 400"
+ARGS+=" --save_best_from=380"
+ARGS+=" --log_frequency=1"
+ARGS+=" --eval_frequency=$EVAL_FREQUENCY"
+ARGS+=" --prediction_frequency=$PREDICTION_FREQUENCY"
+ARGS+=" --grad_accumulation_steps=$GRAD_ACCUMULATION_STEPS "
+ARGS+=" --dali_device=$DALI_DEVICE"
+
+[ "$AMP" = true ] &&                 ARGS+=" --amp"
+[ "$RESUME" = true ] &&              ARGS+=" --resume"
+[ "$CUDNN_BENCHMARK" = true ] &&     ARGS+=" --cudnn_benchmark"
+[ -n "$MAX_DURATION" ] &&            ARGS+=" --override_config input_train.audio_dataset.max_duration=$MAX_DURATION" \
+                                     ARGS+=" --override_config input_train.filterbank_features.max_duration=$MAX_DURATION"
+[ "$PAD_TO_MAX_DURATION" = true ] && ARGS+=" --override_config input_train.audio_dataset.pad_to_max_duration=True" \
+                                     ARGS+=" --override_config input_train.filterbank_features.pad_to_max_duration=True"
+[ -n "$CHECKPOINT" ] &&              ARGS+=" --ckpt=$CHECKPOINT"
+[ -n "$LOG_FILE" ] &&                ARGS+=" --log_file $LOG_FILE"
+[ -n "$PRE_ALLOCATE" ] &&            ARGS+=" --pre_allocate_range $PRE_ALLOCATE"
+
+DISTRIBUTED="-m torch.distributed.launch --nproc_per_node=$NUM_GPUS"
+python $DISTRIBUTED train.py $ARGS
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_performance_1p.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_performance_1p.sh
new file mode 100644
index 0000000000..5183d8282e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_performance_1p.sh
@@ -0,0 +1,95 @@
+#!/bin/bash
+source test/env_npu.sh
+# export ASCEND_SLOG_PRINT_TO_STDOUT=1
+# export ASCEND_GLOBAL_LOG_LEVEL=1
+echo "train log path is 1p_train_performance.log"
+{
+startTime=`date +%Y%m%d-%H:%M`
+startTime_s=`date +%s`
+echo "start time: $startTime"
+export OMP_NUM_THREADS=1
+
+: ${DATA_DIR:=${1:-"/datasets/LibriSpeech"}}
+: ${MODEL_CONFIG:=${2:-"configs/jasper10x5dr_speedp-online_speca.yaml"}}
+: ${OUTPUT_DIR:=${3:-"/results"}}
+: ${CHECKPOINT:=${4:-}}
+: ${RESUME:=true}
+: ${CUDNN_BENCHMARK:=true}
+: ${NUM_GPUS:=1}
+: ${AMP:=True}
+: ${BATCH_SIZE:=32}
+: ${GRAD_ACCUMULATION_STEPS:=2}
+: ${LEARNING_RATE:=0.01}
+: ${MIN_LEARNING_RATE:=0.00001}
+: ${LR_POLICY:=exponential}
+: ${LR_EXP_GAMMA:=0.981}
+: ${EMA:=0.999}
+: ${SEED:=0}
+: ${EPOCHS:=3}
+: ${WARMUP_EPOCHS:=2}
+: ${HOLD_EPOCHS:=140}
+: ${SAVE_FREQUENCY:=10}
+: ${EPOCHS_THIS_JOB:=0}
+: ${DALI_DEVICE:="cpu"}
+: ${PAD_TO_MAX_DURATION:=false}
+: ${EVAL_FREQUENCY:=544}
+: ${PREDICTION_FREQUENCY:=544}
+: ${TRAIN_MANIFESTS:="$DATA_DIR/librispeech-train-clean-100-wav.json \
+                      $DATA_DIR/librispeech-train-clean-360-wav.json \
+                      $DATA_DIR/librispeech-train-other-500-wav.json"}
+: ${VAL_MANIFESTS:="$DATA_DIR/librispeech-dev-clean-wav.json"}
+
+mkdir -p "$OUTPUT_DIR"
+
+ARGS="--dataset_dir=$DATA_DIR"
+ARGS+=" --val_manifests $VAL_MANIFESTS"
+ARGS+=" --train_manifests $TRAIN_MANIFESTS"
+ARGS+=" --model_config=$MODEL_CONFIG"
+ARGS+=" --output_dir=$OUTPUT_DIR"
+ARGS+=" --lr=$LEARNING_RATE"
+ARGS+=" --batch_size=$BATCH_SIZE"
+ARGS+=" --min_lr=$MIN_LEARNING_RATE"
+ARGS+=" --lr_policy=$LR_POLICY"
+ARGS+=" --lr_exp_gamma=$LR_EXP_GAMMA"
+ARGS+=" --epochs=$EPOCHS"
+ARGS+=" --warmup_epochs=$WARMUP_EPOCHS"
+ARGS+=" --hold_epochs=$HOLD_EPOCHS"
+ARGS+=" --epochs_this_job=$EPOCHS_THIS_JOB"
+ARGS+=" --ema=$EMA"
+ARGS+=" --seed=$SEED"
+ARGS+=" --optimizer=novograd"
+ARGS+=" --weight_decay=1e-3"
+ARGS+=" --save_frequency=$SAVE_FREQUENCY"
+ARGS+=" --keep_milestones 100 200 300 400"
+ARGS+=" --save_best_from=380"
+ARGS+=" --log_frequency=1"
+ARGS+=" --eval_frequency=$EVAL_FREQUENCY"
+ARGS+=" --prediction_frequency=$PREDICTION_FREQUENCY"
+ARGS+=" --grad_accumulation_steps=$GRAD_ACCUMULATION_STEPS "
+ARGS+=" --dali_device=$DALI_DEVICE"
+
+[ "$AMP" = true ] &&                 ARGS+=" --amp"
+[ "$RESUME" = true ] &&              ARGS+=" --resume"
+[ "$CUDNN_BENCHMARK" = true ] &&     ARGS+=" --cudnn_benchmark"
+[ -n "$MAX_DURATION" ] &&            ARGS+=" --override_config input_train.audio_dataset.max_duration=$MAX_DURATION" \
+                                     ARGS+=" --override_config input_train.filterbank_features.max_duration=$MAX_DURATION"
+[ "$PAD_TO_MAX_DURATION" = true ] && ARGS+=" --override_config input_train.audio_dataset.pad_to_max_duration=True" \
+                                     ARGS+=" --override_config input_train.filterbank_features.pad_to_max_duration=True"
+[ -n "$CHECKPOINT" ] &&              ARGS+=" --ckpt=$CHECKPOINT"
+[ -n "$LOG_FILE" ] &&                ARGS+=" --log_file $LOG_FILE"
+[ -n "$PRE_ALLOCATE" ] &&            ARGS+=" --pre_allocate_range $PRE_ALLOCATE"
+
+DISTRIBUTED="-m torch.distributed.launch --nproc_per_node=$NUM_GPUS"
+python $DISTRIBUTED train.py $ARGS
+
+
+
+endTime=`date +%Y%m%d-%H:%M`
+endTime_s=`date +%s`
+echo "end time: $endTime"
+sumTime_sec=$(( $endTime_s - $startTime_s ))
+sumTime_min=$(( $sumTime_sec / 60 ))
+echo "$startTime ---> $endTime" "Total: $sumTime_min minutes"
+fps=$(echo $(printf "%.3f" `echo "scale=3; 8707 * 3 / $sumTime_sec"|bc`))
+echo "FPS -----> $fps"
+} > ./1p_train_performance.log 2>&1 & 
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_performance_8p.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_performance_8p.sh
new file mode 100644
index 0000000000..7ccc4018da
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/test/train_performance_8p.sh
@@ -0,0 +1,95 @@
+#!/bin/bash
+source test/env_npu.sh
+# export ASCEND_SLOG_PRINT_TO_STDOUT=1
+# export ASCEND_GLOBAL_LOG_LEVEL=1
+echo "train log path is 8p_train_performance.log"
+{
+startTime=`date +%Y%m%d-%H:%M`
+startTime_s=`date +%s`
+echo "start time: $startTime"
+export OMP_NUM_THREADS=1
+
+: ${DATA_DIR:=${1:-"/datasets/LibriSpeech"}}
+: ${MODEL_CONFIG:=${2:-"configs/jasper10x5dr_speedp-online_speca.yaml"}}
+: ${OUTPUT_DIR:=${3:-"/results"}}
+: ${CHECKPOINT:=${4:-}}
+: ${RESUME:=true}
+: ${CUDNN_BENCHMARK:=true}
+: ${NUM_GPUS:=8}
+: ${AMP:=True}
+: ${BATCH_SIZE:=32}
+: ${GRAD_ACCUMULATION_STEPS:=2}
+: ${LEARNING_RATE:=0.01}
+: ${MIN_LEARNING_RATE:=0.00001}
+: ${LR_POLICY:=exponential}
+: ${LR_EXP_GAMMA:=0.981}
+: ${EMA:=0.999}
+: ${SEED:=0}
+: ${EPOCHS:=3}
+: ${WARMUP_EPOCHS:=2}
+: ${HOLD_EPOCHS:=140}
+: ${SAVE_FREQUENCY:=10}
+: ${EPOCHS_THIS_JOB:=0}
+: ${DALI_DEVICE:="cpu"}
+: ${PAD_TO_MAX_DURATION:=false}
+: ${EVAL_FREQUENCY:=544}
+: ${PREDICTION_FREQUENCY:=544}
+: ${TRAIN_MANIFESTS:="$DATA_DIR/librispeech-train-clean-100-wav.json \
+                      $DATA_DIR/librispeech-train-clean-360-wav.json \
+                      $DATA_DIR/librispeech-train-other-500-wav.json"}
+: ${VAL_MANIFESTS:="$DATA_DIR/librispeech-dev-clean-wav.json"}
+
+mkdir -p "$OUTPUT_DIR"
+
+ARGS="--dataset_dir=$DATA_DIR"
+ARGS+=" --val_manifests $VAL_MANIFESTS"
+ARGS+=" --train_manifests $TRAIN_MANIFESTS"
+ARGS+=" --model_config=$MODEL_CONFIG"
+ARGS+=" --output_dir=$OUTPUT_DIR"
+ARGS+=" --lr=$LEARNING_RATE"
+ARGS+=" --batch_size=$BATCH_SIZE"
+ARGS+=" --min_lr=$MIN_LEARNING_RATE"
+ARGS+=" --lr_policy=$LR_POLICY"
+ARGS+=" --lr_exp_gamma=$LR_EXP_GAMMA"
+ARGS+=" --epochs=$EPOCHS"
+ARGS+=" --warmup_epochs=$WARMUP_EPOCHS"
+ARGS+=" --hold_epochs=$HOLD_EPOCHS"
+ARGS+=" --epochs_this_job=$EPOCHS_THIS_JOB"
+ARGS+=" --ema=$EMA"
+ARGS+=" --seed=$SEED"
+ARGS+=" --optimizer=novograd"
+ARGS+=" --weight_decay=1e-3"
+ARGS+=" --save_frequency=$SAVE_FREQUENCY"
+ARGS+=" --keep_milestones 100 200 300 400"
+ARGS+=" --save_best_from=380"
+ARGS+=" --log_frequency=1"
+ARGS+=" --eval_frequency=$EVAL_FREQUENCY"
+ARGS+=" --prediction_frequency=$PREDICTION_FREQUENCY"
+ARGS+=" --grad_accumulation_steps=$GRAD_ACCUMULATION_STEPS "
+ARGS+=" --dali_device=$DALI_DEVICE"
+
+[ "$AMP" = true ] &&                 ARGS+=" --amp"
+[ "$RESUME" = true ] &&              ARGS+=" --resume"
+[ "$CUDNN_BENCHMARK" = true ] &&     ARGS+=" --cudnn_benchmark"
+[ -n "$MAX_DURATION" ] &&            ARGS+=" --override_config input_train.audio_dataset.max_duration=$MAX_DURATION" \
+                                     ARGS+=" --override_config input_train.filterbank_features.max_duration=$MAX_DURATION"
+[ "$PAD_TO_MAX_DURATION" = true ] && ARGS+=" --override_config input_train.audio_dataset.pad_to_max_duration=True" \
+                                     ARGS+=" --override_config input_train.filterbank_features.pad_to_max_duration=True"
+[ -n "$CHECKPOINT" ] &&              ARGS+=" --ckpt=$CHECKPOINT"
+[ -n "$LOG_FILE" ] &&                ARGS+=" --log_file $LOG_FILE"
+[ -n "$PRE_ALLOCATE" ] &&            ARGS+=" --pre_allocate_range $PRE_ALLOCATE"
+
+DISTRIBUTED="-m torch.distributed.launch --nproc_per_node=$NUM_GPUS"
+python $DISTRIBUTED train.py $ARGS
+
+
+
+endTime=`date +%Y%m%d-%H:%M`
+endTime_s=`date +%s`
+echo "end time: $endTime"
+sumTime_sec=$(( $endTime_s - $startTime_s ))
+sumTime_min=$(( $sumTime_sec / 60 ))
+echo "$startTime ---> $endTime" "Total: $sumTime_min minutes"
+fps=$(echo $(printf "%.3f" `echo "scale=3; 1088 * 3 / $sumTime_sec"|bc`))
+echo "FPS -----> $fps"
+} > ./8p_train_performance.log 2>&1 & 
\ No newline at end of file
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/train.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/train.py
new file mode 100644
index 0000000000..7018852a86
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/train.py
@@ -0,0 +1,528 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+import sklearn
+import argparse
+import copy
+import os
+import random
+import time
+
+try:
+    import nvidia_dlprof_pytorch_nvtx as pyprof
+except ModuleNotFoundError:
+    import pyprof
+
+import torch
+import numpy as np
+import torch.cuda.profiler as profiler
+import torch.distributed as dist
+from apex import amp
+#from apex.parallel import DistributedDataParallel
+from torch.nn.parallel import DistributedDataParallel
+from common import helpers
+# from common.dali.data_loader import DaliDataLoader
+from common.dataset import AudioDataset, get_data_loader
+from common.features import BaseFeatures, FilterbankFeatures
+from common.helpers import (Checkpointer, greedy_wer, num_weights, print_once,
+                            process_evaluation_epoch)
+from common.optimizers import AdamW, lr_policy, Novograd
+from common.tb_dllogger import flush_log, init_log, log
+from common.utils import BenchmarkStats
+from jasper import config
+from jasper.model import CTCLossNM, GreedyCTCDecoder, Jasper
+
+
+def parse_args():
+    parser = argparse.ArgumentParser(description='Jasper')
+
+    training = parser.add_argument_group('training setup')
+    training.add_argument('--epochs', default=400, type=int,
+                          help='Number of epochs for the entire training; influences the lr schedule')
+    training.add_argument("--warmup_epochs", default=0, type=int,
+                          help='Initial epochs of increasing learning rate')
+    training.add_argument("--hold_epochs", default=0, type=int,
+                          help='Constant max learning rate epochs after warmup')
+    training.add_argument('--epochs_this_job', default=0, type=int,
+                          help=('Run for a number of epochs with no effect on the lr schedule.'
+                                'Useful for re-starting the training.'))
+    training.add_argument('--cudnn_benchmark', action='store_true', default=True,
+                          help='Enable cudnn benchmark')
+    training.add_argument('--amp', '--fp16', action='store_true', default=True,
+                          help='Use mixed precision training')
+    training.add_argument('--seed', default=42, type=int, help='Random seed')
+    training.add_argument('--local_rank', default=os.getenv('LOCAL_RANK', 0),
+                          type=int, help='GPU id used for distributed training')
+    training.add_argument('--pre_allocate_range', default=None, type=int, nargs=2,
+                          help='Warmup with batches of length [min, max] before training')
+    training.add_argument('--pyprof', action='store_true', help='Enable pyprof profiling')
+
+    optim = parser.add_argument_group('optimization setup')
+    optim.add_argument('--batch_size', default=32, type=int,
+                       help='Global batch size')
+    optim.add_argument('--lr', default=1e-3, type=float,
+                       help='Peak learning rate')
+    optim.add_argument("--min_lr", default=1e-5, type=float,
+                       help='minimum learning rate')
+    optim.add_argument("--lr_policy", default='exponential', type=str,
+                       choices=['exponential', 'legacy'], help='lr scheduler')
+    optim.add_argument("--lr_exp_gamma", default=0.99, type=float,
+                       help='gamma factor for exponential lr scheduler')
+    optim.add_argument('--weight_decay', default=1e-3, type=float,
+                       help='Weight decay for the optimizer')
+    optim.add_argument('--grad_accumulation_steps', default=1, type=int,
+                       help='Number of accumulation steps')
+    optim.add_argument('--optimizer', default='novograd', type=str,
+                       choices=['novograd', 'adamw'], help='Optimization algorithm')
+    optim.add_argument('--ema', type=float, default=0.0,
+                       help='Discount factor for exp averaging of model weights')
+
+    io = parser.add_argument_group('feature and checkpointing setup')
+    io.add_argument('--dali_device', type=str, choices=['none', 'cpu', 'gpu'],
+                    default='none', help='Use DALI pipeline for fast data processing')
+    io.add_argument('--resume', action='store_true',
+                    help='Try to resume from last saved checkpoint.')
+    io.add_argument('--ckpt', default=None, type=str,
+                    help='Path to a checkpoint for resuming training')
+    io.add_argument('--save_frequency', default=10, type=int,
+                    help='Checkpoint saving frequency in epochs')
+    io.add_argument('--keep_milestones', default=[100, 200, 300], type=int, nargs='+',
+                    help='Milestone checkpoints to keep from removing')
+    io.add_argument('--save_best_from', default=380, type=int,
+                    help='Epoch on which to begin tracking best checkpoint (dev WER)')
+    io.add_argument('--eval_frequency', default=200, type=int,
+                    help='Number of steps between evaluations on dev set')
+    io.add_argument('--log_frequency', default=25, type=int,
+                    help='Number of steps between printing training stats')
+    io.add_argument('--prediction_frequency', default=100, type=int,
+                    help='Number of steps between printing sample decodings')
+    io.add_argument('--model_config', type=str, required=True,
+                    help='Path of the model configuration file')
+    io.add_argument('--train_manifests', type=str, required=True, nargs='+',
+                    help='Paths of the training dataset manifest file')
+    io.add_argument('--val_manifests', type=str, required=True, nargs='+',
+                    help='Paths of the evaluation datasets manifest files')
+    io.add_argument('--dataset_dir', required=True, type=str,
+                    help='Root dir of dataset')
+    io.add_argument('--output_dir', type=str, required=True,
+                    help='Directory for logs and checkpoints')
+    io.add_argument('--log_file', type=str, default=None,
+                    help='Path to save the training logfile.')
+    io.add_argument('--benchmark_epochs_num', type=int, default=1,
+                    help='Number of epochs accounted in final average throughput.')
+    io.add_argument('--override_config', type=str, action='append',
+                    help='Overrides a value from a config .yaml.'
+                         ' Syntax: `--override_config nested.config.key=val`.')
+    return parser.parse_args()
+
+
+def reduce_tensor(tensor, num_gpus):
+    rt = tensor.clone()
+    dist.all_reduce(rt, op=dist.ReduceOp.SUM)
+    return rt.true_divide(num_gpus)
+
+
+def apply_ema(model, ema_model, decay):
+    if not decay:
+        return
+
+    sd = getattr(model, 'module', model).state_dict()
+    for k, v in ema_model.state_dict().items():
+        v.copy_(decay * v + (1 - decay) * sd[k])
+
+
+@torch.no_grad()
+def evaluate(epoch, step, val_loader, val_feat_proc, labels, model,
+             ema_model, ctc_loss, greedy_decoder, use_amp, use_dali=False):
+
+    for model, subset in [(model, 'dev'), (ema_model, 'dev_ema')]:
+        if model is None:
+            continue
+
+        model.eval()
+        start_time = time.time()
+        agg = {'losses': [], 'preds': [], 'txts': []}
+
+        for batch in val_loader:
+            if use_dali:
+                # with DALI, the data is already on GPU
+                feat, feat_lens, txt, txt_lens = batch
+                if val_feat_proc is not None:
+                    feat, feat_lens = val_feat_proc(feat, feat_lens, use_amp)
+            else:
+                audio, audio_lens, txt, txt_lens = batch
+                # print("audio",audio)
+                feat, feat_lens = val_feat_proc(audio, audio_lens, use_amp)
+            feat = feat.npu()
+            audio = audio.npu()
+            feat_lens = feat_lens.npu()
+            txt = txt.npu()
+            log_probs, enc_lens = model.forward(feat, feat_lens)
+            loss = ctc_loss(log_probs, txt, enc_lens, txt_lens)
+            pred = greedy_decoder(log_probs)
+
+            agg['losses'] += helpers.gather_losses([loss])
+            agg['preds'] += helpers.gather_predictions([pred], labels)
+            agg['txts'] += helpers.gather_transcripts([txt], [txt_lens], labels)
+
+        wer, loss = process_evaluation_epoch(agg)
+        log((epoch,), step, subset, {'loss': loss, 'wer': 100.0 * wer,
+                                     'took': time.time() - start_time})
+        model.train()
+    return wer
+
+
+def main():
+    args = parse_args()
+
+    assert(torch.npu.is_available())
+    assert args.prediction_frequency % args.log_frequency == 0
+
+    # torch.backends.cudnn.benchmark = args.cudnn_benchmark
+
+    # set up distributed training
+    multi_gpu = True
+    if multi_gpu:
+        torch.npu.set_device(args.local_rank)
+        dist.init_process_group(backend='hccl', init_method='env://')
+        world_size = dist.get_world_size()
+        print_once(f'Distributed training with {world_size} GPUs\n')
+    else:
+        world_size = 1
+
+    torch.manual_seed(args.seed + args.local_rank)
+    np.random.seed(args.seed + args.local_rank)
+    random.seed(args.seed + args.local_rank)
+
+    init_log(args)
+
+    cfg = config.load(args.model_config)
+    config.apply_config_overrides(cfg, args)
+
+    symbols = helpers.add_ctc_blank(cfg['labels'])
+
+    assert args.grad_accumulation_steps >= 1
+    assert args.batch_size % args.grad_accumulation_steps == 0
+    batch_size = args.batch_size // args.grad_accumulation_steps
+
+    print_once('Setting up datasets...')
+    train_dataset_kw, train_features_kw = config.input(cfg, 'train')
+    val_dataset_kw, val_features_kw = config.input(cfg, 'val')
+
+    # use_dali = args.dali_device in ('cpu', 'gpu')
+    use_dali = False
+    if use_dali:
+        assert train_dataset_kw['ignore_offline_speed_perturbation'], \
+            "DALI doesn't support offline speed perturbation"
+
+        # pad_to_max_duration is not supported by DALI - have simple padders
+        if train_features_kw['pad_to_max_duration']:
+            train_feat_proc = BaseFeatures(
+                pad_align=train_features_kw['pad_align'],
+                pad_to_max_duration=True,
+                max_duration=train_features_kw['max_duration'],
+                sample_rate=train_features_kw['sample_rate'],
+                window_size=train_features_kw['window_size'],
+                window_stride=train_features_kw['window_stride'])
+            train_features_kw['pad_to_max_duration'] = False
+        else:
+            train_feat_proc = None
+
+        if val_features_kw['pad_to_max_duration']:
+            val_feat_proc = BaseFeatures(
+                pad_align=val_features_kw['pad_align'],
+                pad_to_max_duration=True,
+                max_duration=val_features_kw['max_duration'],
+                sample_rate=val_features_kw['sample_rate'],
+                window_size=val_features_kw['window_size'],
+                window_stride=val_features_kw['window_stride'])
+            val_features_kw['pad_to_max_duration'] = False
+        else:
+            val_feat_proc = None
+
+        train_loader = DaliDataLoader(gpu_id=args.local_rank,
+                                      dataset_path=args.dataset_dir,
+                                      config_data=train_dataset_kw,
+                                      config_features=train_features_kw,
+                                      json_names=args.train_manifests,
+                                      batch_size=batch_size,
+                                      grad_accumulation_steps=args.grad_accumulation_steps,
+                                      pipeline_type="train",
+                                      device_type=args.dali_device,
+                                      symbols=symbols)
+
+        val_loader = DaliDataLoader(gpu_id=args.local_rank,
+                                    dataset_path=args.dataset_dir,
+                                    config_data=val_dataset_kw,
+                                    config_features=val_features_kw,
+                                    json_names=args.val_manifests,
+                                    batch_size=batch_size,
+                                    pipeline_type="val",
+                                    device_type=args.dali_device,
+                                    symbols=symbols)
+    else:
+        train_dataset_kw, train_features_kw = config.input(cfg, 'train')
+        train_dataset = AudioDataset(args.dataset_dir,
+                                     args.train_manifests,
+                                     symbols,
+                                     **train_dataset_kw)
+        train_loader = get_data_loader(train_dataset,
+                                       batch_size,
+                                       multi_gpu=multi_gpu,
+                                       shuffle=True,
+                                       num_workers=4)
+        train_feat_proc = FilterbankFeatures(**train_features_kw)
+
+        val_dataset_kw, val_features_kw = config.input(cfg, 'val')
+        val_dataset = AudioDataset(args.dataset_dir,
+                                   args.val_manifests,
+                                   symbols,
+                                   **val_dataset_kw)
+        val_loader = get_data_loader(val_dataset,
+                                     batch_size,
+                                     multi_gpu=multi_gpu,
+                                     shuffle=False,
+                                     num_workers=4,
+                                     drop_last=False)
+        val_feat_proc = FilterbankFeatures(**val_features_kw)
+
+        dur = train_dataset.duration / 3600
+        dur_f = train_dataset.duration_filtered / 3600
+        nsampl = len(train_dataset)
+        print_once(f'Training samples: {nsampl} ({dur:.1f}h, '
+                   f'filtered {dur_f:.1f}h)')
+
+    # if train_feat_proc is not None:
+    #     train_feat_proc.cpu()
+    # if val_feat_proc is not None:
+    #     val_feat_proc.cpu()
+    train_feat_proc.cpu()
+    val_feat_proc.cpu()
+    steps_per_epoch = len(train_loader) // args.grad_accumulation_steps
+
+    # set up the model
+    model = Jasper(encoder_kw=config.encoder(cfg),
+                   decoder_kw=config.decoder(cfg, n_classes=len(symbols)))
+    model.npu()
+    ctc_loss = CTCLossNM(n_classes=len(symbols))
+    greedy_decoder = GreedyCTCDecoder()
+
+    print_once(f'Model size: {num_weights(model) / 10**6:.1f}M params\n')
+
+    # optimization
+    kw = {'lr': args.lr, 'weight_decay': args.weight_decay}
+    if args.optimizer == "novograd":
+        optimizer = Novograd(model.parameters(), **kw)
+    elif args.optimizer == "adamw":
+        optimizer = AdamW(model.parameters(), **kw)
+    else:
+        raise ValueError(f'Invalid optimizer "{args.optimizer}"')
+
+    adjust_lr = lambda step, epoch, optimizer: lr_policy(
+        step, epoch, args.lr, optimizer, steps_per_epoch=steps_per_epoch,
+        warmup_epochs=args.warmup_epochs, hold_epochs=args.hold_epochs,
+        num_epochs=args.epochs, policy=args.lr_policy, min_lr=args.min_lr,
+        exp_gamma=args.lr_exp_gamma)
+
+    if args.amp:
+#        model, optimizer = amp.initialize(
+#            min_loss_scale=1.0, models=model, optimizers=optimizer,
+#            opt_level='O1', max_loss_scale=512.0)
+         model, optimizer = amp.initialize(models=model, optimizers=optimizer, loss_scale=32)
+
+    if args.ema > 0:
+        ema_model = copy.deepcopy(model)
+    else:
+        ema_model = None
+
+    if multi_gpu:
+        model = DistributedDataParallel(model, device_ids=[args.local_rank], broadcast_buffers=False)
+
+    if args.pyprof:
+        pyprof.init(enable_function_stack=True)
+
+    # load checkpoint
+    meta = {'best_wer': 10**6, 'start_epoch': 0}
+    checkpointer = Checkpointer(args.output_dir, 'Jasper',
+                                args.keep_milestones, args.amp)
+    if args.resume:
+        args.ckpt = checkpointer.last_checkpoint() or args.ckpt
+
+    if args.ckpt is not None:
+        checkpointer.load(args.ckpt, model, ema_model, optimizer, meta)
+
+    start_epoch = meta['start_epoch']
+    best_wer = meta['best_wer']
+    epoch = 1
+    step = start_epoch * steps_per_epoch + 1
+
+    if args.pyprof:
+        torch.autograd.profiler.emit_nvtx().__enter__()
+        profiler.start()
+
+    # training loop
+    model.train()
+
+    # pre-allocate
+    if args.pre_allocate_range is not None:
+        n_feats = train_features_kw['n_filt']
+        pad_align = train_features_kw['pad_align']
+        a, b = args.pre_allocate_range
+        for n_frames in range(a, b + pad_align, pad_align):
+            print_once(f'Pre-allocation ({batch_size}x{n_feats}x{n_frames})...')
+
+            feat = torch.randn(batch_size, n_feats, n_frames, device='cpu')
+            feat_lens = torch.ones(batch_size, device='cpu').fill_(n_frames)
+            txt = torch.randint(high=len(symbols)-1, size=(batch_size, 100),
+                                device='cpu')
+            txt_lens = torch.ones(batch_size, device='cpu').fill_(100)
+            log_probs, enc_lens = model(feat, feat_lens)
+            del feat
+            loss = ctc_loss(log_probs, txt, enc_lens, txt_lens)
+            loss.backward()
+            model.zero_grad()
+
+    bmark_stats = BenchmarkStats()
+
+    for epoch in range(start_epoch + 1, args.epochs + 1):
+        if multi_gpu and not use_dali:
+            train_loader.sampler.set_epoch(epoch)
+
+        epoch_utts = 0
+        epoch_loss = 0
+        accumulated_batches = 0
+        epoch_start_time = time.time()
+
+        for batch in train_loader:
+
+            if accumulated_batches == 0:
+                adjust_lr(step, epoch, optimizer)
+                optimizer.zero_grad()
+                step_loss = 0
+                step_utts = 0
+                step_start_time = time.time()
+
+            if use_dali:
+                # with DALI, the data is already on GPU
+                feat, feat_lens, txt, txt_lens = batch
+                if train_feat_proc is not None:
+                    feat, feat_lens = train_feat_proc(feat, feat_lens, args.amp)
+            else:
+                # print("progress")
+                # batch = [t.cpu(non_blocking=True) for t in batch]
+                audio, audio_lens, txt, txt_lens = batch
+                # print("audio",audio)
+                feat, feat_lens = train_feat_proc(audio, audio_lens, args.amp)
+                # print("feat",feat)
+                # print("feat_len",feat_lens)
+                feat = feat.npu()
+                audio = audio.npu()
+                feat_lens = feat_lens.npu()
+                txt = txt.npu()
+            log_probs, enc_lens = model(feat, feat_lens)
+
+            loss = ctc_loss(log_probs, txt, enc_lens, txt_lens)
+            loss /= args.grad_accumulation_steps
+
+            if torch.isnan(loss).any():
+                print_once(f'WARNING: loss is NaN; skipping update')
+            else:
+                if multi_gpu:
+                    step_loss += reduce_tensor(loss.data, world_size).item()
+                else:
+                    step_loss += loss.item()
+
+                if args.amp:
+                    with amp.scale_loss(loss, optimizer) as scaled_loss:
+                        scaled_loss.backward()
+                else:
+                    loss.backward()
+                step_utts += batch[0].size(0) * world_size
+                epoch_utts += batch[0].size(0) * world_size
+                accumulated_batches += 1
+
+            if accumulated_batches % args.grad_accumulation_steps == 0:
+                epoch_loss += step_loss
+                optimizer.step()
+                apply_ema(model, ema_model, args.ema)
+
+                if step % args.log_frequency == 0:
+                    preds = greedy_decoder(log_probs)
+                    wer, pred_utt, ref = greedy_wer(preds, txt, txt_lens, symbols)
+
+                    if step % args.prediction_frequency == 0:
+                        print_once(f'  Decoded:   {pred_utt[:90]}')
+                        print_once(f'  Reference: {ref[:90]}')
+
+                    step_time = time.time() - step_start_time
+                    log((epoch, step % steps_per_epoch or steps_per_epoch, steps_per_epoch),
+                        step, 'train',
+                        {'loss': step_loss,
+                         'wer': 100.0 * wer,
+                         'throughput': step_utts / step_time,
+                         'took': step_time,
+                         'lrate': optimizer.param_groups[0]['lr']})
+
+                step_start_time = time.time()
+
+                if step % args.eval_frequency == 0:
+                    wer = evaluate(epoch, step, val_loader, val_feat_proc,
+                                   symbols, model, ema_model, ctc_loss,
+                                   greedy_decoder, args.amp, use_dali)
+
+                    if wer < best_wer and epoch >= args.save_best_from:
+                        checkpointer.save(model, ema_model, optimizer, epoch,
+                                          step, best_wer, is_best=True)
+                        best_wer = wer
+
+                step += 1
+                accumulated_batches = 0
+                # end of step
+
+            # DALI iterator need to be exhausted;
+            # if not using DALI, simulate drop_last=True with grad accumulation
+            if not use_dali and step > steps_per_epoch * epoch:
+                break
+
+        epoch_time = time.time() - epoch_start_time
+        epoch_loss /= steps_per_epoch
+        log((epoch,), None, 'train_avg', {'throughput': epoch_utts / epoch_time,
+                                          'took': epoch_time,
+                                          'loss': epoch_loss})
+        bmark_stats.update(epoch_utts, epoch_time, epoch_loss)
+
+        if epoch % args.save_frequency == 0 or epoch in args.keep_milestones:
+            checkpointer.save(model, ema_model, optimizer, epoch, step, best_wer)
+
+        if 0 < args.epochs_this_job <= epoch - start_epoch:
+            print_once(f'Finished after {args.epochs_this_job} epochs.')
+            break
+        # end of epoch
+
+    if args.pyprof:
+        profiler.stop()
+        torch.autograd.profiler.emit_nvtx().__exit__(None, None, None)
+
+    log((), None, 'train_avg', bmark_stats.get(args.benchmark_epochs_num))
+
+    if epoch == args.epochs:
+        evaluate(epoch, step, val_loader, val_feat_proc, symbols, model,
+                 ema_model, ctc_loss, greedy_decoder, args.amp, use_dali)
+
+        checkpointer.save(model, ema_model, optimizer, epoch, step, best_wer)
+    flush_log()
+
+
+if __name__ == "__main__":
+    main()
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/Dockerfile b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/Dockerfile
new file mode 100644
index 0000000000..9fda344c5c
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/Dockerfile
@@ -0,0 +1,10 @@
+ARG FROM_IMAGE_NAME=nvcr.io/nvidia/tritonserver:20.10-py3-clientsdk
+FROM ${FROM_IMAGE_NAME}
+
+RUN apt update && apt install -y python3-pyaudio libsndfile1
+
+RUN pip3 install -U pip
+RUN pip3 install onnxruntime unidecode inflect soundfile
+
+WORKDIR /workspace/jasper
+COPY . .
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/README.md b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/README.md
new file mode 100644
index 0000000000..38dd13ec2d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/README.md
@@ -0,0 +1,388 @@
+# Deploying the Jasper Inference model using Triton Inference Server
+
+This subfolder of the Jasper for PyTorch repository contains scripts for  deployment of high-performance inference on NVIDIA Triton Inference Server as well as detailed performance analysis. It offers different options for the inference model pipeline.
+
+
+## Table Of Contents
+- [Solution overview](#solution-overview)
+- [Inference Pipeline in Triton Inference Server](#inference-pipeline-in-triton-inference-server)
+- [Setup](#setup)
+- [Quick Start Guide](#quick-start-guide)
+- [Advanced](#advanced)
+     * [Scripts and sample code](#scripts-and-sample-code)
+- [Performance](#performance)
+     * [Inference Benchmarking in Triton Inference Server](#inference-benchmarking-in-triton-inference-server)
+     * [Results](#results)
+       * [Performance Analysis for Triton Inference Server: NVIDIA T4](#performance-analysis-for-triton-inference-server-nvidia-t4)
+       * [Maximum batch size](#maximum-batch-size)
+            * [Batching techniques: Static versus Dynamic Batching](#batching-techniques-static-versus-dynamic)
+            * [TensorRT, ONNXRT-CUDA, and PyTorch JIT comparisons](#tensorrt-onnxrt-cuda-and-pytorch-jit-comparisons)
+- [Release Notes](#release-notes)
+	* [Changelog](#change-log)
+	* [Known issues](#known-issues)
+
+
+## Solution Overview
+
+The [NVIDIA Triton Inference Server](https://github.com/NVIDIA/triton-inference-server) provides a datacenter and cloud inferencing solution optimized for NVIDIA GPUs. The server provides an inference service via an HTTP or gRPC endpoint, allowing remote clients to request inferencing for any number of GPU or CPU models being managed by the server.
+
+This folder contains detailed performance analysis as well as scripts to run Jasper inference using Triton Inference Server.
+
+A typical Triton Inference Server pipeline can be broken down into the following steps:
+
+1. The client serializes the inference request into a message and sends it to the server (Client Send).
+2. The message travels over the network from the client to the server (Network).
+3. The message arrives at the server, and is deserialized (Server Receive).
+4. The request is placed on the queue (Server Queue).
+5. The request is removed from the queue and computed (Server Compute).
+6. The completed request is serialized in a message and sent back to the client (Server Send).
+7. The completed message then travels over the network from the server to the client (Network).
+8. The completed message is deserialized by the client and processed as a completed inference request (Client Receive).
+
+Generally, for local clients, steps 1-4 and 6-8 will only occupy a small fraction of time, compared to step 5. As backend deep learning systems like Jasper are rarely exposed directly to end users, but instead only interfacing with local front-end servers, for the sake of Jasper, we can consider that all clients are local.
+
+In this section, we will go over how to launch both the Triton Inference Server and the client and get the best performance solution that fits your specific application needs.
+
+More information on how to perform inference using NVIDIA Triton Inference Server can be found in [triton/README.md](https://github.com/triton-inference-server/server/blob/master/README.md).
+
+
+## Inference Pipeline in Triton Inference Server
+
+The Jasper model pipeline consists of 3 components, where each part can be customized to be a different backend:
+
+**Data preprocessor**
+
+The data processor transforms an input raw audio file into a spectrogram. By default the pipeline uses mel filter banks as spectrogram features. This part does not have any learnable weights.
+
+**Acoustic model**
+
+The acoustic model takes in the spectrogram and outputs a probability over a list of characters. This part is the most compute intensive, taking more than 90% of the entire end-to-end pipeline. The acoustic model is the only component with learnable parameters and what differentiates Jasper from other end-to-end neural speech recognition models. In the original paper, the acoustic model contains a masking operation for training (More details in [Jasper PyTorch README](../README.md)). We do not use masking for inference.
+
+**Greedy decoder**
+
+The decoder takes the probabilities over the list of characters and outputs the final transcription. Greedy decoding is a fast and simple way of doing this by always choosing the character with the maximum probability.
+
+To run a model with TensorRT, we first construct the model in PyTorch, which is then exported into a ONNX static graph. Finally, a TensorRT engine is constructed from the ONNX file and can be launched to do inference. The following table shows which backends are supported for each part along the model pipeline.
+
+|Backend\Pipeline component|Data preprocessor|Acoustic Model|Decoder|
+|---|---|---|---|
+|PyTorch JIT|x|x|x|
+|ONNX|-|x|-|
+|TensorRT|-|x|-|
+
+In order to run inference with TensorRT outside of the inference server, refer to the [Jasper TensorRT README](../tensort/README.md).
+
+
+## Setup
+
+The repository contains a folder `./triton` with a `Dockerfile` which extends the PyTorch 20.10-py3 NGC container and encapsulates some dependencies. Ensure you have the following components:
+
+- [NVIDIA Docker](https://github.com/NVIDIA/nvidia-docker)
+- [PyTorch 20.10-py3 NGC container](https://ngc.nvidia.com/catalog/containers/nvidia:pytorch)
+- [Triton Inference Server 20.10 NGC container](https://ngc.nvidia.com/catalog/containers/nvidia:tritonserver)
+- Access to [NVIDIA machine learning repository](https://developer.download.nvidia.com/compute/machine-learning/repos/ubuntu1804/x86_64/nvidia-machine-learning-repo-ubuntu1804_1.0.0-1_amd64.deb) and [NVIDIA CUDA repository](https://developer.download.nvidia.com/compute/cuda/repos/ubuntu1804/x86_64/cuda-repo-ubuntu1804_10.1.243-1_amd64.deb) for NVIDIA TensorRT 6
+- Supported GPUs:
+    - [NVIDIA Volta architecture](https://www.nvidia.com/en-us/data-center/volta-gpu-architecture/)
+    - [NVIDIA Turing architecture](https://www.nvidia.com/en-us/geforce/turing/)
+    - [NVIDIA Ampere architecture](https://www.nvidia.com/en-us/data-center/nvidia-ampere-gpu-architecture/)
+- [Pretrained Jasper Model Checkpoint](https://ngc.nvidia.com/catalog/models/nvidia:jasper_pyt_ckpt_amp)
+
+Required Python packages are listed in `requirements.txt`. These packages are automatically installed when the Docker container is built.
+
+
+## Quick Start Guide
+
+Running the following scripts will build and launch the container containing all required dependencies for native PyTorch as well as Triton. This is necessary for using inference and can also be used for data download, processing, and training of the model. For more information on the scripts and arguments, refer to the [Advanced](#advanced) section.
+
+1. Clone the repository.
+
+    ```bash
+    git clone https://github.com/NVIDIA/DeepLearningExamples
+    cd DeepLearningExamples/PyTorch/SpeechRecognition/Jasper
+    ```
+
+2. Build the Jasper PyTorch container.
+
+    Running the following scripts will build the container which contains all the required dependencies for data download and processing as well as converting the model.
+
+    ```bash
+    bash scripts/docker/build.sh
+    ```
+
+3. Start an interactive session in the Docker container:
+
+    ```bash
+    bash scripts/docker/launch.sh <DATA_DIR> <CHECKPOINT_DIR> <RESULT_DIR>
+    ```
+
+    Where <DATA_DIR>, <CHECKPOINT_DIR> and <RESULT_DIR> can be either empty or absolute directory paths to dataset, existing checkpoints or potential output files. When left empty, they default to `datasets/`, `/checkpoints`, and `results/`, respectively. The `/datasets`, `/checkpoints`, `/results` directories will be mounted as volumes and mapped to the corresponding directories `<DATA_DIR>`, `<CHECKPOINT_DIR>`, `<RESULT_DIR>` on the host.
+
+    Note that `<DATA_DIR>`, `<CHECKPOINT_DIR>`, and `<RESULT_DIR>` directly correspond to the same arguments in `scripts/docker/launch.sh` and `trt/scripts/docker/launch.sh` mentioned in the [Jasper PyTorch README](../README.md) and [Jasper TensorRT README](../tensorrt/README.md).
+
+    Briefly, `<DATA_DIR>` should contain, or be prepared to contain a `LibriSpeech` sub-directory (created in [Acquiring Dataset](../trt/README.md)), `<CHECKPOINT_DIR>` should contain a PyTorch model checkpoint (`*.pt`) file obtained through training described in [Jasper PyTorch README](../README.md), and `<RESULT_DIR>` should be prepared to contain converted model and logs.
+
+4. Downloading the `test-clean` part of `LibriSpeech` is required for model conversion. But it is not required for inference on Triton Inference Server, which can use a single .wav audio file. To download and preprocess LibriSpeech, run the following inside the container:
+
+   ```bash
+   bash triton/scripts/download_triton_librispeech.sh
+   bash triton/scripts/preprocess_triton_librispeech.sh
+   ```
+
+5. (Option 1) Convert pretrained PyTorch model checkpoint into Triton Inference Server compatible model backends.
+
+   Inside the container, run:
+
+    ```bash
+    export CHECKPOINT_PATH=<CHECKPOINT_PATH>
+    export CONVERT_PRECISIONS=<CONVERT_PRECISIONS>
+    export CONVERTS=<CONVERTS>
+    bash triton/scripts/export_model.sh
+    ```
+
+   Where `<CHECKPOINT_PATH>` (`"/checkpoints/jasper_fp16.pt"`) is the absolute file path of the pretrained checkpoint, `<CONVERT_PRECISIONS>` (`"fp16" "fp32"`) is the list of precisions used for conversion, and `<CONVERTS>` (`"feature-extractor" "decoder" "ts-trace" "onnx" "tensorrt"`) is the list of conversions to be applied. The feature extractor converts only to TorchScript trace module (`feature-extractor`), the decoder only to TorchScript script module (`decoder`), and the Jasper model can convert to TorchScript trace module (`ts-trace`), ONNX (`onnx`), or TensorRT (`tensorrt`).
+
+   A pretrained PyTorch model checkpoint for model conversion can be downloaded from the [NGC model repository](https://ngc.nvidia.com/catalog/models/nvidia:jasper_pyt_ckpt_amp).
+
+   More details can be found in the [Advanced](#advanced) section under [Scripts and sample code](#scripts-and-sample-code).
+
+6. (Option 2) Download pre-exported inference checkpoints from NGC.
+
+   Alternatively, you can skip the manual model export and download already generated model backends for every version of the model pipeline.
+
+   * [Jasper_ONNX](https://ngc.nvidia.com/catalog/models/nvidia:jasper_pyt_onnx_fp16_amp/version),
+   * [Jasper_TorchScript](https://ngc.nvidia.com/catalog/models/nvidia:jasper_pyt_torchscript_fp16_amp/version),
+   * [Jasper_TensorRT_Turing](https://ngc.nvidia.com/catalog/models/nvidia:jasper_pyt_trt_fp16_amp_turing/version),
+   * [Jasper_TensorRT_Volta](https://ngc.nvidia.com/catalog/models/nvidia:jasper_pyt_trt_fp16_amp_volta/version).
+
+   If you wish to use TensorRT pipeline, make sure to download the correct version for your hardware. The extracted model folder should contain 3 subfolders `feature-extractor-ts-trace`, `decoder-ts-script` and `jasper-x` where `x` can be `ts-trace`, `onnx`, `tensorrt` depending on the model backend. Copy the 3 model folders to the directory `./triton/model_repo/fp16` in your Jasper project.
+
+7. Build a container that extends Triton Inference Client:
+
+    From outside the container, run:
+
+    ```bash
+    bash triton/scripts/docker/build_triton_client.sh
+    ```
+
+Once the above steps are completed you can either run inference benchmarks or perform inference on real data.
+
+8. (Option 1) Run all inference benchmarks.
+
+    From outside the container, run:
+
+    ```bash
+    export RESULT_DIR=<RESULT_DIR>
+    export PRECISION_TESTS=<PRECISION_TESTS>
+    export BATCH_SIZES=<BATCH_SIZES>
+    export SEQ_LENS=<SEQ_LENS>
+    bash triton/scripts/execute_all_perf_runs.sh
+    ```
+
+    Where `<RESULT_DIR>` is the absolute path to potential output files (`./results`), `<PRECISION_TESTS>` is a list of precisions to be tested (`"fp16" "fp32"`), `<BATCH_SIZES>` is a list of tested batch sizes (`"1" "2" "4" "8"`), and `<SEQ_LENS>` are tested sequnce lengths (`"32000" "112000" "267200"`).
+
+    Note: This can take several hours to complete due to the extensiveness of the benchmark. More details about the benchmark are found in the [Advanced](#advanced) section under [Performance](#performance).
+
+9. (Option 2) Run inference on real data using the Client and Triton Inference Server.
+
+    8.1 From outside the container, restart the server:
+
+    ```bash
+    bash triton/scripts/run_server.sh <MODEL_TYPE> <PRECISION>
+    ```
+
+    8.2 From outside the container, submit the client request using:
+    ```bash
+    bash triton/scripts/run_client.sh <MODEL_TYPE> <DATA_DIR> <FILE>
+    ```
+
+    Where `<MODEL_TYPE>` can be either "ts-trace", "tensorrt" or "onnx", `<PRECISION>` is either "fp32" or "fp16". `<DATA_DIR>` is an absolute local path to the directory of files. <FILE> is the relative path to <DATA_DIR> to either an audio file in .wav format or a manifest file in .json format.
+
+    Note: If <FILE> is *.json <DATA_DIR> should be the path to the LibriSpeech dataset. In this case this script will do both inference and evaluation on the accoring LibriSpeech dataset.
+
+
+## Advanced
+
+The following sections provide greater details about the Triton Inference Server pipeline and inference analysis and benchmarking results.
+
+
+### Scripts and sample code
+
+The `triton/` directory contains the following files:
+* `jasper-client.py`: Python client script that takes an audio file and a specific model pipeline type and submits a client request to the server to run inference with the model on the given audio file.
+* `speech_utils.py`: helper functions for `jasper-client.py`.
+* `converter.py`: Python script for model conversion to different backends.
+* `jasper_module.py`: helper functions for `converter.py`.
+* `model_repo_configs/`: directory with Triton model config files for different backend and precision configurations.
+
+The `triton/scripts/` directory has easy to use scripts to run supported functionalities, such as:
+* `./docker/build_triton_client.sh`: builds container
+* `execute_all_perf_runs.sh`: runs all benchmarks using Triton Inference Server performance client; calls `generate_perf_results.sh`
+* `export_model.sh`: from pretrained PyTorch checkpoint generates backends for every version of the model inference pipeline.
+* `prepare_model_repository.sh`: copies model config files from `./model_repo_configs/` to `./deploy/model_repo` and creates links to generated model backends, setting up the model repository for Triton Inference Server
+* `generate_perf_results.sh`: runs benchmark with `perf-client` for specific configuration and calls `run_perf_client.sh`
+* `run_server.sh`: launches Triton Inference Server
+* `run_client.sh`: launches client by using `jasper-client.py` to submit inference requests to server
+
+
+### Running the Triton Inference Server
+
+Launch the Triton Inference Server in detached mode to run in the background by default:
+
+```bash
+bash triton/scripts/run_server.sh
+```
+
+To run in the foreground interactively, for debugging purposes, run:
+
+```bash
+DAEMON="--detach=false" bash triton/scripts/run_server.sh
+```
+
+The script mounts and loads models at `$PWD/triton/deploy/model_repo` to the server with all visible GPUs. In order to selectively choose the devices, set `NVIDIA_VISIBLE_DEVICES`.
+
+
+### Running the Triton Inference Client
+
+*Real data*
+In order to run the client with real data, run:
+
+```bash
+bash triton/scripts/run_client.sh <backend> <data directory> <audio file>
+```
+
+The script calls `triton/jasper-client.py` which preprocesses data and sends/receives requests to/from the server.
+
+*Synthetic data*
+In order to run the client with synthetic data for performance measurements, run:
+
+```bash
+export MODEL_NAME=jasper-tensorrt-ensemble
+export MODEL_VERSION=1
+export BATCH_SIZE=1
+export MAX_LATENCY=500
+export MAX_CONCURRENCY=64
+export AUDIO_LENGTH=32000
+export SERVER_HOSTNAME=localhost
+export RESULT_DIR_H=${PWD}/results/perf_client/${MODEL_NAME}/batch_${BATCH_SIZE}_len_${AUDIO_LENGTH}
+bash triton/scripts/run_perf_client.sh
+```
+
+The export values above are default values. The script waits until the server is up and running, sends requests as per the constraints set and writes results to `/results/results_${TIMESTAMP}.csv` where `TIMESTAMP=$(date "+%y%m%d_%H%M")` and `/results/` is the results directory mounted in the docker .
+
+For more information about `perf_client`, refer to the [official documentation](https://docs.nvidia.com/deeplearning/triton-inference-server/master-user-guide/docs/optimization.html#perf-client).
+
+
+## Performance
+
+The performance measurements in this document were conducted at the time of publication and may not reflect the performance achieved from NVIDIA’s latest software release. For the most up-to-date performance measurements, go to [NVIDIA Data Center Deep Learning Product Performance](https://developer.nvidia.com/deep-learning-performance-training-inference).
+
+### Inference Benchmarking in Triton Inference Server
+
+To benchmark the inference performance on Volta Turing or Ampere GPU, run `bash triton/scripts/execute_all_perf_runs.sh` according to [Quick-Start-Guide](#quick-start-guide) Step 7.
+
+By default, this script measures inference performance for all 3 model pipelines: PyTorch JIT  (`ts-trace`) pipeline, ONNX (`onnx`) pipeline, TensorRT(`tensorrt`) pipeline, both with FP32 and FP16 precision. Each of these pipelines is measured for different audio input lengths (2sec, 7sec, 16.7sec) and a range of different server batch sizes (up to 8). This takes place in `triton/scripts/generate_perf_results.sh`. For a specific audio length and batch size, static and dynamic batching comparison is performed.
+
+
+### Results
+
+In the following section, we analyze the results using the example of the Triton pipeline.
+
+#### Performance Analysis for Triton Inference Server: NVIDIA T4
+
+All results below are obtained using the following configurations:
+* Single T4 16GB GPU on a local server
+* FP16 precision
+* Python 3.6.10
+* PyTorch 1.7.0a0+7036e91
+* TensorRT 7.2.1.4
+* CUDA 11.1.0.024
+* CUDNN 8.0.4.30
+
+##### Batching techniques: Static Batching
+
+Static batching is a feature of the inference server that allows inference requests to be served as they are received. The largest improvements to throughput come from increasing the batch size due to efficiency gains in the GPU with larger batches.
+
+![](../images/static_fp16_2s.png)
+Figure 1: Throughput vs. Latency for Jasper, Audio Length = 2sec using various model backends available in Triton Inference Server and static batching.
+
+![](../images/static_fp16_7s.png)
+Figure 2: Throughput vs. Latency for Jasper, Audio Length = 7sec using various model backends available in Triton Inference Server and static batching.
+
+![](../images/static_fp16_16.7s.png)
+Figure 3: Throughput vs. Latency for Jasper, Audio Length = 16.7sec using various model backends available in Triton Inference Server and static batching.
+
+These charts can be used to establish the optimal batch size to use in dynamic batching, given a latency budget. For example, in Figure 2 (Audio length = 7s) given a budget of 50ms, the optimal batch size to use for the TensorRT backend is 4. This will result in a maximum throughput of 100 inf/s under the latency constraint. In all three charts, TensorRT shows the best throughput and latency performance for a given batch size
+
+##### Batching techniques: Dynamic Batching
+
+Dynamic batching is a feature of the inference server that allows inference requests to be combined by the server, so that a batch is created dynamically, resulting in an increased throughput. It is preferred in scenarios where we would like to maximize throughput and GPU utilization at the cost of higher latencies. You can set the Dynamic Batcher parameter `max_queue_delay_microseconds` to indicate the maximum amount of time you are willing to wait and `preferred_batch_size` to indicate your maximum server batch size in the Triton Inference Server model config.
+
+Figures 4, 5, and 6 emphasizes the increase in overall throughput with dynamic batching. At low numbers of concurrent requests, the increased throughput comes at the cost of increasing latency as the requests are queued up to max_queue_delay_microseconds.
+
+![](../images/tensorrt_2s.png)
+Figure 4: Triton pipeline - Latency & Throughput vs Concurrency using dynamic Batching at maximum server batch size = 8, max_queue_delay_microseconds = 5000, input audio length = 2 seconds, TensorRT backend.
+
+![](../images/tensorrt_7s.png)
+Figure 5: Triton pipeline - Latency & Throughput vs Concurrency using dynamic Batching at maximum server batch size =  8, max_queue_delay_microseconds = 5000, input audio length = 7 seconds, TensorRT backend.
+
+![](../images/tensorrt_16.7s.png)
+Figure 6: Triton pipeline - Latency & Throughput vs Concurrency using dynamic Batching at maximum server batch size = 8, max_queue_delay_microseconds = 5000, input audio length = 16.7 seconds, TensorRT backend.
+
+##### TensorRT, ONNXRT-CUDA, and PyTorch JIT comparisons
+
+The following tables show inference and latency comparisons across all 3 backends for mixed precision and static batching. The main observations are:
+Increasing the batch size leads to higher inference throughput and - latency up to a certain batch size, after which it slowly saturates.
+The longer the audio length, the lower the throughput and the higher the latency.
+
+###### Throughput Comparison
+
+The following table shows the throughput benchmark results for all 3 model backends in Triton Inference Server using static batching under optimal concurrency
+
+|Audio length in seconds|Batch Size|TensorRT (inf/s)|PyTorch (inf/s)|ONNXRT-CUDA (inf/s)|TensorRT/PyTorch Speedup|TensorRT/ONNXRT-CUDA Speedup|
+|---    |---    |---    |---    |---    |---    |---    |
+|  2.0| 1|  49.67|  55.67|  41.67| 0.89| 1.19|
+|  2.0| 2|  98.67|  96.00|  77.33| 1.03| 1.28|
+|  2.0| 4| 180.00| 141.33| 118.67| 1.27| 1.52|
+|  2.0| 8| 285.33| 202.67| 136.00| 1.41| 2.10|
+|  7.0| 1|  47.67|  37.00|  18.00| 1.29| 2.65|
+|  7.0| 2|  79.33|  47.33|  46.00| 1.68| 1.72|
+|  7.0| 4| 100.00|  73.33|  36.00| 1.36| 2.78|
+|  7.0| 8| 117.33|  82.67|  40.00| 1.42| 2.93|
+| 16.7| 1|  36.33|  21.67|  11.33| 1.68| 3.21|
+| 16.7| 2|  40.67|  25.33|  16.00| 1.61| 2.54|
+| 16.7| 4|  46.67|  37.33|  16.00| 1.25| 2.92|
+| 16.7| 8|  48.00|  40.00|  18.67| 1.20| 2.57|
+
+###### Latency Comparison
+
+The following table shows the throughput benchmark results for all 3 model backends in Triton Inference Server using static batching and a single concurrent request.
+
+|Audio length in seconds|Batch Size|TensorRT (ms)|PyTorch (ms)|ONNXRT-CUDA (ms)|TensorRT/PyTorch Speedup|TensorRT/ONNXRT-CUDA Speedup|
+|---    |---    |---    |---    |---    |---    |---    |
+|  2.0| 1|  23.61|  25.06| 31.84| 1.06| 1.35|
+|  2.0| 2|  24.56|  25.11| 37.54| 1.02| 1.53|
+|  2.0| 4|  25.90|  31.00| 37.20| 1.20| 1.44|
+|  2.0| 8|  31.57|  41.76| 37.13| 1.32| 1.18|
+|  7.0| 1|  24.79|  30.55| 32.16| 1.23| 1.30|
+|  7.0| 2|  28.48|  45.05| 37.47| 1.58| 1.32|
+|  7.0| 4|  41.71|  57.71| 37.92| 1.38| 0.91|
+|  7.0| 8|  72.19|  98.84| 38.13| 1.37| 0.53|
+| 16.7| 1|  30.66|  48.42| 32.74| 1.58| 1.07|
+| 16.7| 2|  52.79|  81.89| 37.82| 1.55| 0.72|
+| 16.7| 4|  92.86| 115.03| 37.91| 1.24| 0.41|
+| 16.7| 8| 170.34| 203.52| 37.84| 2.36| 0.22|
+
+## Release Notes
+
+### Changelog
+
+March 2021
+* Updated ONNX runtime information
+
+February 2021
+* Updated Triton scripts for compatibility with Triton Inference Server version 2
+* Updated Quick Start Guide
+* Updated performance results
+
+### Known issues
+There are no known issues in this deployment.
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/converter.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/converter.py
new file mode 100644
index 0000000000..9286551ca4
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/converter.py
@@ -0,0 +1,254 @@
+# *****************************************************************************
+#  Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#      * Redistributions of source code must retain the above copyright
+#        notice, this list of conditions and the following disclaimer.
+#      * Redistributions in binary form must reproduce the above copyright
+#        notice, this list of conditions and the following disclaimer in the
+#        documentation and/or other materials provided with the distribution.
+#      * Neither the name of the NVIDIA CORPORATION nor the
+#        names of its contributors may be used to endorse or promote products
+#        derived from this software without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#  DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************
+
+
+import os
+import json
+import torch
+import argparse
+import importlib
+from pytorch.utils import extract_io_props, load_io_props
+
+import logging
+
+def get_parser():
+
+    parser = argparse.ArgumentParser()
+    # required args
+    parser.add_argument("--model-module", type=str, default="", required=True,
+                        help="Module with model initializer and data loader")
+    parser.add_argument('--convert', choices=['ts-script', 'ts-trace',
+                                              'onnx', 'tensorrt'],
+                        required=True, help='convert to '
+                        'ts-script: TorchScript using torch.jit.script, '
+                        'ts-trace: TorchScript using torch.jit.trace, '
+                        'onnx: ONNX using torch.onnx.export, '
+                        'tensorrt: TensorRT using OnnxParser, ')
+    parser.add_argument("--max_workspace_size", type=int,
+                        default=512*1024*1024,
+                        help="set the size of the workspace for TensorRT \
+                        conversion")
+    parser.add_argument("--precision", choices=['fp16', 'fp32'],
+                        default='fp32', help="convert TensorRT or \
+                        TorchScript model in a given precision")
+    parser.add_argument('--convert-filename', type=str, default='',
+                        help='Saved model name')
+    parser.add_argument('--save-dir', type=str, default='',
+                        help='Saved model directory')
+    parser.add_argument("--max-batch-size", type=int, default=1,
+                        help="Specifies the 'max_batch_size' in the Triton \
+                           model config and in TensorRT builder. See the \
+                           Triton and TensorRT documentation for more info.")
+    parser.add_argument('--device', type=str, default='cuda',
+                        help='Select device for conversion.')
+
+    parser.add_argument('model_arguments', nargs=argparse.REMAINDER,
+                        help='arguments that will be ignored by \
+                           converter lib and will be forwarded to your convert \
+                           script')
+
+    return parser
+
+
+class Converter:
+    def __init__(self, model, dataloader):
+        self.model = model
+        self.dataloader = dataloader
+
+        self.convert_props = {
+            'ts-script': {
+                'convert_func':  self.to_torchscript,
+                'convert_filename': 'model.pt'
+            },
+            'ts-trace': {
+                'convert_func'    :  self.to_torchscript,
+                'convert_filename': 'model.pt'
+            },
+            'onnx': {
+                'convert_func'    :  self.to_onnx,
+                'convert_filename': 'model.onnx'
+            },
+            'tensorrt': {
+                'convert_func'    :  self.to_tensorrt,
+                'convert_filename': 'model.plan'
+            }
+        }
+
+    def convert(self, convert_type, save_dir, convert_filename,
+                device, precision='fp32',
+                max_batch_size=1,
+                # args for TensorRT:
+                max_workspace_size=None):
+        ''' convert the model '''
+        self.convert_type = convert_type
+        self.max_workspace_size = max_workspace_size
+        self.max_batch_size = max_batch_size
+        self.precision = precision
+
+        # override default name if user provided name
+        if convert_filename != '':
+            self.convert_props[convert_type]['convert_filename'] = convert_filename
+
+        # setup device
+        torch_device = torch.device(device)
+
+        # prepare model
+        self.model.to(torch_device)
+        self.model.eval()
+        assert (not self.model.training), \
+            "[Converter error]: could not set the model to eval() mode!"
+
+        io_props = None
+        if self.dataloader is not None:
+            io_props = extract_io_props(self.model, self.dataloader,
+                                        torch_device, precision, max_batch_size)
+
+        assert self.convert_type == "ts-script" or io_props is not None, \
+            "Input and output properties are empty. For conversion types \
+other than \'ts-script\' input shapes are required to generate dummy input. \
+Make sure that dataloader works correctly or that IO props file is provided."
+
+        # prepare save path
+        model_name = self.convert_props[convert_type]['convert_filename']
+        convert_model_path = os.path.join(save_dir, model_name)
+
+        # get convert method depending on the convert type
+        convert_func = self.convert_props[convert_type]['convert_func']
+
+        # convert the model - will be saved to disk
+        if self.convert_type == "tensorrt":
+            io_filepath = "triton/tensorrt_io_props_" + str(precision) + ".json"
+            io_props = load_io_props(io_filepath)
+
+        convert_func(model, torch_device, io_props, convert_model_path)
+
+        assert (os.path.isfile(convert_model_path)), \
+            f"[Converter error]: saving model to {convert_model_path} failed!"
+
+
+    def generate_dummy_input(self, io_props, device):
+
+        from pytorch.utils import triton_type_to_torch_type
+
+        dummy_input = []
+        for s,t in zip(io_props['opt_shapes'], io_props['input_types']):
+            t = triton_type_to_torch_type[t]
+            tensor = torch.empty(size=s, dtype=t, device=device).random_()
+            dummy_input.append(tensor)
+        dummy_input = tuple(dummy_input)
+
+        return dummy_input
+
+    def to_onnx(self, model, device, io_props, convert_model_path):
+        ''' convert the model to onnx '''
+
+        dummy_input = self.generate_dummy_input(io_props, device)
+
+        opset_version = 11
+        # convert the model to onnx
+        with torch.no_grad():
+            torch.onnx.export(model, dummy_input,
+                              convert_model_path,
+                              do_constant_folding=True,
+                              input_names=io_props['input_names'],
+                              output_names=io_props['output_names'],
+                              dynamic_axes=io_props['dynamic_axes'],
+                              opset_version=opset_version,
+                              enable_onnx_checker=True)
+
+
+    def to_tensorrt(self, model, device, io_props, convert_model_path):
+        ''' convert the model to tensorrt '''
+
+        assert (self.max_workspace_size), "[Converter error]: for TensorRT conversion you must provide \'max_workspace_size\'."
+
+        import tensorrt as trt
+        from pytorch.utils import build_tensorrt_engine
+
+        # convert the model to onnx first
+        self.to_onnx(model, device, io_props, convert_model_path)
+        del model
+        torch.cuda.empty_cache()
+
+        zipped = zip(io_props['input_names'], io_props['min_shapes'],
+                     io_props['opt_shapes'], io_props['max_shapes'])
+        shapes = []
+        for name,min_shape,opt_shape,max_shape in zipped:
+            d = {"name":name, "min": min_shape,
+                 "opt": opt_shape, "max": max_shape}
+            shapes.append(d)
+
+        tensorrt_fp16 = True if self.precision == 'fp16' else False
+        # build tensorrt engine
+        engine = build_tensorrt_engine(convert_model_path, shapes,
+                                       self.max_workspace_size,
+                                       self.max_batch_size,
+                                       tensorrt_fp16)
+
+        assert engine is not None, "[Converter error]: TensorRT build failure"
+
+        # write tensorrt engine
+        with open(convert_model_path, 'wb') as f:
+            f.write(engine.serialize())
+
+
+    def to_torchscript(self, model, device, io_props, convert_model_path):
+        ''' convert the model to torchscript '''
+
+        if self.convert_type == 'ts-script':
+
+            model_ts = torch.jit.script(model)
+
+        else: # self.convert_type == 'ts-trace'
+
+            dummy_input = self.generate_dummy_input(io_props, device)
+            with torch.no_grad():
+                model_ts = torch.jit.trace(model, dummy_input)
+
+        # save the model
+        torch.jit.save(model_ts, convert_model_path)
+
+
+if __name__=='__main__':
+
+    parser = get_parser()
+    args = parser.parse_args()
+    model_args_list = args.model_arguments[1:]
+
+    logging.basicConfig(level=logging.INFO)
+
+    mm = importlib.import_module(args.model_module)
+
+    model = mm.init_model(model_args_list, args.precision, args.device)
+
+    dataloader = mm.get_dataloader(model_args_list)
+    converter = Converter(model, dataloader)
+
+    converter.convert(args.convert, args.save_dir, args.convert_filename,
+                      args.device, args.precision,
+                      args.max_batch_size,
+                      args.max_workspace_size)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/jasper-client.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/jasper-client.py
new file mode 100644
index 0000000000..c1a556b04a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/jasper-client.py
@@ -0,0 +1,400 @@
+#!/usr/bin/python
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+
+import sys
+import argparse
+import numpy as np
+import os
+from speech_utils import AudioSegment, SpeechClient
+import soundfile
+import pyaudio as pa
+import threading
+import math
+import time
+import glob
+
+FLAGS = None
+
+
+# read audio chunk from a file
+def get_audio_chunk_from_soundfile(sf, chunk_size, int_values):
+
+    dtype = 'int32' if int_values else 'float32'
+    audio_signal = sf.read(chunk_size, dtype=dtype)
+    end = False
+    # pad to chunk size
+    if len(audio_signal) < chunk_size:
+        end = True
+        audio_signal = np.pad(audio_signal, (0, chunk_size-len(
+            audio_signal)), mode='constant')
+    return audio_signal, end
+
+
+# generator that returns chunks of audio data from file
+def audio_generator_from_file(input_filename, target_sr, int_values,
+                              chunk_duration):
+
+    sf = soundfile.SoundFile(input_filename, 'rb')
+    chunk_size = int(chunk_duration*sf.samplerate)
+    start = True
+    end = False
+
+    while not end:
+
+        audio_signal, end = get_audio_chunk_from_soundfile(
+            sf, chunk_size, int_values)
+
+        audio_segment = AudioSegment(audio_signal, sf.samplerate, target_sr)
+
+        yield audio_segment.samples, target_sr, start, end
+
+        start = False
+
+    sf.close()
+
+
+# generator that returns chunks of audio data from file
+class AudioGeneratorFromMicrophone:
+
+    def __init__(self,input_device_id, target_sr, chunk_duration):
+
+        self.recording_state = "init"
+        self.target_sr  = target_sr
+        self.chunk_duration = chunk_duration
+
+        self.p = pa.PyAudio()
+
+        device_info = self.p.get_host_api_info_by_index(0)
+        num_devices = device_info.get('deviceCount')
+        devices = {}
+        for i in range(0, num_devices):
+            if (self.p.get_device_info_by_host_api_device_index(0, i).get(
+                'maxInputChannels')) > 0:
+                devices[i] = self.p.get_device_info_by_host_api_device_index(
+                    0, i)
+
+        if (len(devices) == 0):
+            raise RuntimeError("Cannot find any valid input devices")
+
+        if input_device_id is None or input_device_id not in \
+            devices.keys():
+            print("\nInput Devices:")
+            for id, info in devices.items():
+                print("{}: {}".format(id,info.get("name")))
+            input_device_id = int(input("Enter device id to use: "))
+
+        self.input_device_id = input_device_id
+
+
+    def generate_audio(self):
+
+        chunk_size = int(self.chunk_duration*self.target_sr)
+
+
+        self. recording_state = "init"
+
+        def keyboard_listener():
+            input("Press Enter to start and end recording...")
+            self.recording_state = "capture"
+            print("Recording...")
+
+            input("")
+            self.recording_state = "release"
+
+        listener = threading.Thread(target=keyboard_listener)
+        listener.start()
+
+        start = True
+        end = False
+
+        stream_initialized = False
+        step = 0
+        while self.recording_state != "release":
+            try:
+                if self.recording_state == "capture":
+
+                    if not stream_initialized:
+                        stream = self.p.open(
+                            format=pa.paInt16,
+                            channels=1,
+                            rate=self.target_sr,
+                            input=True,
+                            input_device_index=self.input_device_id,
+                            frames_per_buffer=chunk_size)
+                        stream_initialized = True
+
+                    # Read audio chunk from microphone
+                    audio_signal = stream.read(chunk_size)
+                    audio_signal = np.frombuffer(audio_signal,dtype=np.int16)
+                    audio_segment = AudioSegment(audio_signal,
+                                                              self.target_sr,
+                                                              self.target_sr)
+
+                    yield audio_segment.samples, self.target_sr, start, end
+
+                    start = False
+                    step += 1
+            except Exception as e:
+                print(e)
+                break
+
+        stream.close()
+        self.p.terminate()
+
+    def generate_audio_signal(self):
+
+
+        #chunk_size = int(self.chunk_duration*self.target_sr)
+        chunk_size = int(0.2*self.target_sr)
+        self. recording_state = "init"
+
+        def keyboard_listener():
+            input("Press Enter to start and end recording...")
+            self.recording_state = "capture"
+            print("Recording...")
+
+            input("")
+            self.recording_state = "release"
+
+        listener = threading.Thread(target=keyboard_listener)
+        listener.start()
+
+        audio_samples = []
+        stream_initialized = False
+        step = 0
+        while self.recording_state != "release":
+            try:
+                if self.recording_state == "capture":
+
+                    if not stream_initialized:
+                        stream = self.p.open(
+                            format=pa.paInt16,
+                            channels=1,
+                            rate=self.target_sr,
+                            input=True,
+                            input_device_index=self.input_device_id,
+                            frames_per_buffer=chunk_size)
+                        stream_initialized = True
+
+                    # Read audio chunk from microphone
+                    audio_signal = stream.read(chunk_size)
+                    audio_signal = np.frombuffer(audio_signal,dtype=np.int16)
+                    audio_segment = AudioSegment(audio_signal,
+                                                              self.target_sr,
+                                                              self.target_sr)
+
+                    if step == 0:
+                        audio_samples = audio_segment.samples
+                    else:
+                        audio_samples = np.concatenate((audio_samples,
+                                                       audio_segment.samples))
+
+                    start = False
+                    step += 1
+            except Exception as e:
+                print(e)
+                break
+
+        stream.close()
+        self.p.terminate()
+
+        return audio_samples
+
+# generator that returns chunks of audio features from file
+def audio_features_generator(input_filename, speech_features_params,
+                             target_sr, int_values, chunk_duration):
+
+    sf = soundfile.SoundFile(input_filename, 'rb')
+
+    chunk_size = int(chunk_duration*sf.samplerate)
+
+    start = True
+    end = False
+
+    while not end:
+
+        audio_signal, end = get_audio_chunk_from_soundfile(sf, chunk_size,
+                                                       int_values)
+
+        audio_segment = AudioSegment(audio_signal, sf.samplerate, target_sr)
+        audio_features, features_length = get_speech_features(
+          audio_segment.samples, target_sr, speech_features_params)
+
+        yield audio_features, start, end
+
+        start = False
+
+    sf.close()
+
+
+def audio_features_generator_with_buffer(input_filename,
+                                         speech_features_params, target_sr,
+                                         int_values, chunk_duration):
+
+    sf = soundfile.SoundFile(input_filename, 'rb')
+
+    chunk_size = int(chunk_duration*sf.samplerate)
+
+    start = True
+    end = False
+
+    audio_signal = np.zeros(shape=3*chunk_size, dtype=np.float32)
+
+    while not end:
+
+        audio_signal[-chunk_size:], end = get_audio_chunk_from_soundfile(sf, chunk_size, int_values)
+
+        audio_segment = AudioSegment(audio_signal, sf.samplerate, target_sr)
+        audio_features, features_length = get_speech_features(
+          audio_segment.samples, target_sr, speech_features_params)
+
+        yield audio_features, start, end
+
+        start = False
+        audio_signal[:-chunk_size] = audio_signal[chunk_size:]
+
+
+    sf.close()
+
+
+if __name__ == '__main__':
+    parser = argparse.ArgumentParser()
+    parser.add_argument('-v', '--verbose', action="store_true", required=False,
+                        default=False, help='Enable verbose output')
+    parser.add_argument('--fixed_size', type=int, required=False,
+                        default=0,
+                        help="send fixed_size requests, pad or truncate")    
+    parser.add_argument('--batch_size', type=int, required=False, default=1,
+                        help='batch size')
+    parser.add_argument('--model_platform', required=False,
+                        default='trt',
+                        help='Jasper model platform')
+    parser.add_argument('-u', '--url', type=str, required=False,
+                        default='localhost:8000',
+                        help='Inference server URL. Default is '
+                             'localhost:8000.')
+    parser.add_argument('-i', '--protocol', type=str, required=False,
+                        default='HTTP',
+                        help='Protocol (HTTP/gRPC) used to communicate with '
+                             'inference service. Default is HTTP.')
+    parser.add_argument('--audio_filename', type=str, required=False,
+                        default=None,
+                        help='Input audio filename')
+    parser.add_argument('--data_dir', type=str, required=False,
+                        default=None,
+                        help='data directory')
+    parser.add_argument('--manifest_filename', type=str, required=False,
+                        default=None,
+                        help='relative manifest paths to --data_dir directory.')
+
+    FLAGS = parser.parse_args()
+
+    protocol = FLAGS.protocol.lower()
+
+    valid_model_platforms = {"ts-trace","onnx", "tensorrt"}
+
+    if FLAGS.model_platform not in valid_model_platforms:
+        raise ValueError("Invalid model_platform {}. Valid choices are {"
+                         "}".format(FLAGS.model_platform,
+            valid_model_platforms))
+
+
+    model_name = "jasper-" + FLAGS.model_platform + "-ensemble"
+
+    speech_client = SpeechClient(
+        FLAGS.url, protocol, model_name, 1,
+        FLAGS.batch_size, model_platform=FLAGS.model_platform,
+        verbose=FLAGS.verbose, mode="synchronous",
+        from_features=False
+    )
+    
+    filenames = []
+    transcripts = []
+    if FLAGS.audio_filename is not None:
+        audio_file = os.path.join(FLAGS.data_dir, FLAGS.audio_filename)
+        if os.path.isdir(audio_file):
+            filenames = glob.glob(os.path.join(os.path.abspath(audio_file), "**", "*.wav"),
+                                recursive=True)
+        else:
+            filenames = [audio_file]
+    elif FLAGS.manifest_filename is not None:
+        filter_speed=1.0
+        data_dir=FLAGS.data_dir
+        labels = [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'", "<BLANK>"]
+        labels_map = dict([(labels[i], i) for i in range(len(labels))])
+        blank_index = len(labels)-1
+        table = None
+        import string
+        punctuation = string.punctuation
+        punctuation = punctuation.replace("+", "")
+        punctuation = punctuation.replace("&", "")
+        table = str.maketrans(punctuation, " " * len(punctuation))
+
+        import json
+        if "./triton" not in sys.path:
+            sys.path.append("./")
+            sys.path.append("./triton")
+        from speech_utils import normalize_string, parse_transcript
+        FLAGS.manifest_filename = FLAGS.manifest_filename.split(',')
+        for manifest in FLAGS.manifest_filename:
+            manifest=os.path.join(data_dir, manifest)
+            print(manifest)
+            with open(manifest, "r", encoding="utf-8") as fh:
+                a=json.load(fh)
+                for data in a:
+                    files_and_speeds = data['files']
+                    audio_path = [x['fname'] for x in files_and_speeds if x['speed'] == filter_speed][0]
+                    filenames.append(os.path.join(data_dir, audio_path))
+                    transcript_text = data['transcript']
+                    transcript_text = normalize_string(transcript_text, labels=labels, table=table)
+                    transcripts.append(transcript_text) #parse_transcript(transcript_text, labels_map, blank_index)) # convert to vocab indices
+
+        
+        
+    # Read the audio files
+    # Group requests in batches
+    audio_idx = 0
+    last_request = False
+    predictions = []
+    while not last_request:
+        batch_audio_samples = []
+        batch_filenames = []
+        
+        for idx in range(FLAGS.batch_size):
+            filename = filenames[audio_idx]
+            print("Reading audio file: ", filename)
+            audio = AudioSegment.from_file(
+                filename,
+                offset=0, duration=FLAGS.fixed_size).samples
+            if FLAGS.fixed_size:
+                audio = np.resize(audio, FLAGS.fixed_size)
+                
+            audio_idx = (audio_idx + 1) % len(filenames)
+            if audio_idx == 0:
+                last_request = True
+
+            batch_audio_samples.append(audio)
+            batch_filenames.append(filename)
+
+        predictions += speech_client.recognize(
+            batch_audio_samples,
+            batch_filenames)
+
+    if transcripts:
+        predictions = [x for l in predictions for x in l ]
+        from metrics import word_error_rate
+        wer, scores, num_words = word_error_rate(predictions, transcripts)
+        print(wer)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/jasper_module.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/jasper_module.py
new file mode 100644
index 0000000000..5636aacf83
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/jasper_module.py
@@ -0,0 +1,174 @@
+# *****************************************************************************
+#  Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#      * Redistributions of source code must retain the above copyright
+#        notice, this list of conditions and the following disclaimer.
+#      * Redistributions in binary form must reproduce the above copyright
+#        notice, this list of conditions and the following disclaimer in the
+#        documentation and/or other materials provided with the distribution.
+#      * Neither the name of the NVIDIA CORPORATION nor the
+#        names of its contributors may be used to endorse or promote products
+#        derived from this software without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#  DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************
+
+import torch
+import sys
+sys.path.append("./")
+
+class FeatureCollate:
+
+    def __init__(self, feature_proc):
+        self.feature_proc = feature_proc
+
+    def __call__(self, batch):
+        bs = len(batch)
+        max_len = lambda l,idx: max(el[idx].size(0) for el in l)
+        audio = torch.zeros(bs, max_len(batch, 0))
+        audio_lens = torch.zeros(bs, dtype=torch.int32)
+
+        for i, sample in enumerate(batch):
+            audio[i].narrow(0, 0, sample[0].size(0)).copy_(sample[0])
+            audio_lens[i] = sample[1]
+
+        ret = (audio, audio_lens)
+
+        if self.feature_proc is not None:
+            feats, feat_lens = self.feature_proc(audio, audio_lens)
+            ret = (feats,)
+
+        return ret
+
+
+def get_dataloader(model_args_list):
+    ''' return dataloader for inference '''
+
+    from inference import get_parser
+    from common.helpers import add_ctc_blank
+    from jasper import config
+    from common.dataset import (AudioDataset, FilelistDataset, get_data_loader,
+                                SingleAudioDataset)
+    from common.features import FilterbankFeatures
+
+    parser = get_parser()
+    parser.add_argument('--component', type=str, default="model",
+                        choices=["feature-extractor", "model", "decoder"],
+                        help='Component to convert')
+    args = parser.parse_args(model_args_list)
+
+
+    if args.component == "decoder":
+        return None
+
+    cfg = config.load(args.model_config)
+    config.apply_config_overrides(cfg, args)
+
+    symbols = add_ctc_blank(cfg['labels'])
+
+    dataset_kw, features_kw = config.input(cfg, 'val')
+
+    dataset = AudioDataset(args.dataset_dir, args.val_manifests,
+                           symbols, **dataset_kw)
+
+    data_loader = get_data_loader(dataset, args.batch_size, multi_gpu=False,
+                                  shuffle=False, num_workers=4, drop_last=False)
+    feature_proc = None
+
+    if args.component == "model":
+        feature_proc = FilterbankFeatures(**features_kw)
+
+    data_loader.collate_fn = FeatureCollate(feature_proc)
+
+    return data_loader
+
+
+def init_feature_extractor(args):
+
+    from jasper import config
+    from common.features import FilterbankFeatures
+
+    cfg = config.load(args.model_config)
+    config.apply_config_overrides(cfg, args)
+    _, features_kw = config.input(cfg, 'val')
+
+    feature_proc = FilterbankFeatures(**features_kw)
+
+    return feature_proc
+
+
+def init_acoustic_model(args):
+
+    from common.helpers import add_ctc_blank
+    from jasper.model import Jasper
+    from jasper import config
+
+    cfg = config.load(args.model_config)
+    config.apply_config_overrides(cfg, args)
+
+    if cfg['jasper']['encoder']['use_conv_masks'] == True:
+        print("[Jasper module]: Warning: setting 'use_conv_masks' \
+to False; masked convolutions are not supported.")
+        cfg['jasper']['encoder']['use_conv_masks'] = False
+
+    symbols = add_ctc_blank(cfg['labels'])
+    model = Jasper(encoder_kw=config.encoder(cfg),
+                   decoder_kw=config.decoder(cfg, n_classes=len(symbols)))
+
+    if args.ckpt is not None:
+        checkpoint = torch.load(args.ckpt, map_location="cpu")
+        key = 'ema_state_dict' if args.ema else 'state_dict'
+        state_dict = checkpoint[key]
+        model.load_state_dict(state_dict, strict=True)
+
+    return model
+
+
+def init_decoder(args):
+
+    class GreedyCTCDecoderSimple(torch.nn.Module):
+        @torch.no_grad()
+        def forward(self, log_probs):
+            return log_probs.argmax(dim=-1, keepdim=False).int()
+    return GreedyCTCDecoderSimple()
+
+
+def init_model(model_args_list, precision, device):
+    ''' Return either of the components: feature-extractor, model, or decoder.
+The returned compoenent is ready to convert '''
+
+    from inference import get_parser
+    parser = get_parser()
+    parser.add_argument('--component', type=str, default="model",
+                        choices=["feature-extractor", "model", "decoder"],
+                        help='Component to convert')
+    args = parser.parse_args(model_args_list)
+
+    init_comp = {"feature-extractor": init_feature_extractor,
+                 "model": init_acoustic_model,
+                 "decoder": init_decoder}
+    comp = init_comp[args.component](args)
+
+    torch_device = torch.device(device)
+    print(f"[Jasper module]: using device {torch_device}")
+    comp.to(torch_device)
+    comp.eval()
+
+    if precision == "fp16":
+        print("[Jasper module]: using mixed precision")
+        comp.half()
+    else:
+        print("[Jasper module]: using fp32 precision")
+    return comp
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/decoder-ts-script/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/decoder-ts-script/config.pbtxt
new file mode 100644
index 0000000000..3cd610b8ec
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/decoder-ts-script/config.pbtxt
@@ -0,0 +1,46 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+name: "decoder-ts-script"
+platform: "pytorch_libtorch"
+default_model_filename: "model.pt"
+max_batch_size: 64
+
+input [
+    {
+	name: "input__0"
+	data_type: TYPE_FP16
+	dims: [ -1, 29 ]
+    }
+]
+
+output [
+    {
+	name: "output__0"
+	data_type: TYPE_INT32
+	dims: [-1]
+    }
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/feature-extractor-ts-trace/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/feature-extractor-ts-trace/config.pbtxt
new file mode 100644
index 0000000000..12b92764a0
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/feature-extractor-ts-trace/config.pbtxt
@@ -0,0 +1,32 @@
+name: "feature-extractor-ts-trace"
+platform: "pytorch_libtorch"
+default_model_filename: "model.pt"
+max_batch_size: 64
+
+input [
+    {
+	name: "input__0"
+	data_type: TYPE_FP16
+	dims: [ -1 ]
+    },
+    {
+	name: "input__1"
+	data_type: TYPE_INT32
+	dims: [ 1 ]
+	reshape { shape: [] }
+    }
+]
+
+output [
+    {
+	name: "output__0"
+	data_type: TYPE_FP16
+	dims: [64, -1]
+    },
+    {	
+	name: "output__1"
+	data_type: TYPE_INT32
+	dims: [ 1 ]
+	reshape: { shape: [] }
+    }
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-onnx-ensemble/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-onnx-ensemble/config.pbtxt
new file mode 100644
index 0000000000..e058b1ee1e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-onnx-ensemble/config.pbtxt
@@ -0,0 +1,63 @@
+name: "jasper-onnx-ensemble"
+platform: "ensemble"
+max_batch_size: 8#MAX_BATCH
+input {
+    name: "AUDIO_SIGNAL"
+    data_type: TYPE_FP16
+    dims: -1#AUDIO_LENGTH
+}
+input {
+    name: "NUM_SAMPLES"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
+}
+output {
+    name: "TRANSCRIPT"
+    data_type: TYPE_INT32
+    dims: [-1]
+}
+ensemble_scheduling {
+    step {
+	model_name: "feature-extractor-ts-trace"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_SIGNAL"
+	}
+	input_map {
+	    key: "input__1"
+	    value: "NUM_SAMPLES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "AUDIO_FEATURES"
+	}    
+    }
+    step {
+	model_name: "jasper-onnx"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_FEATURES"
+
+	}
+	output_map {
+	    key: "output__0"
+	    value: "CHARACTER_PROBABILITIES"
+
+	}
+
+    }
+    step {
+	model_name: "decoder-ts-script"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "CHARACTER_PROBABILITIES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "TRANSCRIPT"
+	}
+    }
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-onnx/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-onnx/config.pbtxt
new file mode 100644
index 0000000000..037c3eb556
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-onnx/config.pbtxt
@@ -0,0 +1,58 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+name: "jasper-onnx"
+platform: "onnxruntime_onnx"
+default_model_filename: "model.onnx"
+
+max_batch_size: 8#MAX_BATCH
+
+input [
+    {
+        name: "input__0"
+        data_type: TYPE_FP16
+        dims: [64, -1]
+    }
+]
+
+output [
+    {
+        name: "output__0"
+        data_type: TYPE_FP16
+        dims: [-1, 29 ]
+    }
+]
+
+instance_group {
+  count: 1#NUM_ENGINES
+  gpus: 0
+  kind: KIND_GPU
+}
+
+#db#dynamic_batching {
+#db#    preferred_batch_size: 8#MAX_BATCH
+#db#    max_queue_delay_microseconds: #MAX_QUEUE
+#db#}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-tensorrt-ensemble/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-tensorrt-ensemble/config.pbtxt
new file mode 100644
index 0000000000..7a0b897435
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-tensorrt-ensemble/config.pbtxt
@@ -0,0 +1,63 @@
+name: "jasper-tensorrt-ensemble"
+platform: "ensemble"
+max_batch_size: 8#MAX_BATCH
+input {
+    name: "AUDIO_SIGNAL"
+    data_type: TYPE_FP16
+    dims: -1#AUDIO_LENGTH
+}
+input {
+    name: "NUM_SAMPLES"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
+}
+output {
+    name: "TRANSCRIPT"
+    data_type: TYPE_INT32
+    dims: [-1]
+}
+ensemble_scheduling {
+    step {
+	model_name: "feature-extractor-ts-trace"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_SIGNAL"
+	}
+	input_map {
+	    key: "input__1"
+	    value: "NUM_SAMPLES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "AUDIO_FEATURES"
+	}    
+    }
+    step {
+	model_name: "jasper-tensorrt"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_FEATURES"
+
+	}
+	output_map {
+	    key: "output__0"
+	    value: "CHARACTER_PROBABILITIES"
+
+	}
+
+    }
+    step {
+	model_name: "decoder-ts-script"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "CHARACTER_PROBABILITIES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "TRANSCRIPT"
+	}
+    }
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-tensorrt/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-tensorrt/config.pbtxt
new file mode 100644
index 0000000000..1d8f133f8d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-tensorrt/config.pbtxt
@@ -0,0 +1,58 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+name: "jasper-tensorrt"
+platform: "tensorrt_plan"
+default_model_filename: "model.plan"
+
+max_batch_size: 8#MAX_BATCH
+
+input [
+    {
+        name: "input__0"
+        data_type: TYPE_FP16
+        dims: [64, -1]
+    }      
+]
+
+output [
+    {
+        name: "output__0"
+        data_type: TYPE_FP16
+        dims: [-1, 29 ]
+    }
+]
+
+instance_group {
+  count: 1#NUM_ENGINES
+  gpus: 0
+  kind: KIND_GPU
+}
+
+#db#dynamic_batching {
+#db#    preferred_batch_size: 8#MAX_BATCH
+#db#    max_queue_delay_microseconds: #MAX_QUEUE
+#db#}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-ts-trace-ensemble/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-ts-trace-ensemble/config.pbtxt
new file mode 100644
index 0000000000..87d3969047
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-ts-trace-ensemble/config.pbtxt
@@ -0,0 +1,63 @@
+name: "jasper-ts-trace-ensemble"
+platform: "ensemble"
+max_batch_size: 8#MAX_BATCH
+input {
+    name: "AUDIO_SIGNAL"
+    data_type: TYPE_FP16
+    dims: -1#AUDIO_LENGTH
+}
+input {
+    name: "NUM_SAMPLES"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
+}
+output {
+    name: "TRANSCRIPT"
+    data_type: TYPE_INT32
+    dims: [-1]
+}
+ensemble_scheduling {
+    step {
+	model_name: "feature-extractor-ts-trace"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_SIGNAL"
+	}
+	input_map {
+	    key: "input__1"
+	    value: "NUM_SAMPLES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "AUDIO_FEATURES"
+	}    
+    }
+    step {
+	model_name: "jasper-ts-trace"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_FEATURES"
+
+	}
+	output_map {
+	    key: "output__0"
+	    value: "CHARACTER_PROBABILITIES"
+
+	}
+
+    }
+    step {
+	model_name: "decoder-ts-script"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "CHARACTER_PROBABILITIES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "TRANSCRIPT"
+	}
+    }
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-ts-trace/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-ts-trace/config.pbtxt
new file mode 100644
index 0000000000..56c14df4fe
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp16/jasper-ts-trace/config.pbtxt
@@ -0,0 +1,32 @@
+name: "jasper-ts-trace"
+platform: "pytorch_libtorch"
+default_model_filename: "model.pt"
+
+max_batch_size: 8#MAX_BATCH
+
+input [
+  {
+    name: "input__0"
+    data_type: TYPE_FP16
+    dims: [64, -1]
+  }
+]
+
+output [
+  {
+    name: "output__0"
+    data_type: TYPE_FP16
+    dims: [-1, 29]
+  }
+]
+
+instance_group {
+  count: 1#NUM_ENGINES
+  gpus: 0
+  kind: KIND_GPU
+}
+
+#db#dynamic_batching {
+#db#    preferred_batch_size: 8#MAX_BATCH
+#db#    max_queue_delay_microseconds: #MAX_QUEUE
+#db#}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/decoder-ts-script/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/decoder-ts-script/config.pbtxt
new file mode 100644
index 0000000000..dcefb06f1a
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/decoder-ts-script/config.pbtxt
@@ -0,0 +1,46 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+name: "decoder-ts-script"
+platform: "pytorch_libtorch"
+default_model_filename: "model.pt"
+max_batch_size: 64
+
+input [
+    {
+	name: "input__0"
+	data_type: TYPE_FP32
+	dims: [ -1, 29 ]
+    }
+]
+
+output [
+    {
+	name: "output__0"
+	data_type: TYPE_INT32
+	dims: [-1]
+    }
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/feature-extractor-ts-trace/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/feature-extractor-ts-trace/config.pbtxt
new file mode 100644
index 0000000000..cb29295f8e
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/feature-extractor-ts-trace/config.pbtxt
@@ -0,0 +1,32 @@
+name: "feature-extractor-ts-trace"
+platform: "pytorch_libtorch"
+default_model_filename: "model.pt"
+max_batch_size: 64
+
+input [
+    {
+	name: "input__0"
+	data_type: TYPE_FP32
+	dims: [ -1 ]
+    },
+    {
+	name: "input__1"
+	data_type: TYPE_INT32
+	dims: [ 1 ]
+	reshape { shape: [] }
+    }
+]
+
+output [
+    {
+	name: "output__0"
+	data_type: TYPE_FP32
+	dims: [64, -1]
+    },
+    {	
+	name: "output__1"
+	data_type: TYPE_INT32
+	dims: [ 1 ]
+	reshape: { shape: [] }
+    }
+]
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-onnx-ensemble/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-onnx-ensemble/config.pbtxt
new file mode 100644
index 0000000000..96a96f3e3f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-onnx-ensemble/config.pbtxt
@@ -0,0 +1,63 @@
+name: "jasper-onnx-ensemble"
+platform: "ensemble"
+max_batch_size: 8#MAX_BATCH
+input {
+    name: "AUDIO_SIGNAL"
+    data_type: TYPE_FP32
+    dims: -1#AUDIO_LENGTH
+}
+input {
+    name: "NUM_SAMPLES"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
+}
+output {
+    name: "TRANSCRIPT"
+    data_type: TYPE_INT32
+    dims: [-1]
+}
+ensemble_scheduling {
+    step {
+	model_name: "feature-extractor-ts-trace"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_SIGNAL"
+	}
+	input_map {
+	    key: "input__1"
+	    value: "NUM_SAMPLES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "AUDIO_FEATURES"
+	}    
+    }
+    step {
+	model_name: "jasper-onnx"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_FEATURES"
+
+	}
+	output_map {
+	    key: "output__0"
+	    value: "CHARACTER_PROBABILITIES"
+
+	}
+
+    }
+    step {
+	model_name: "decoder-ts-script"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "CHARACTER_PROBABILITIES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "TRANSCRIPT"
+	}
+    }
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-onnx/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-onnx/config.pbtxt
new file mode 100644
index 0000000000..66690f88ac
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-onnx/config.pbtxt
@@ -0,0 +1,58 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+name: "jasper-onnx"
+platform: "onnxruntime_onnx"
+default_model_filename: "model.onnx"
+
+max_batch_size: 8#MAX_BATCH
+
+input [
+    {
+        name: "input__0"
+        data_type: TYPE_FP32
+        dims: [64, -1]
+    }      
+]
+
+output [
+    {
+        name: "output__0"
+        data_type: TYPE_FP32
+        dims: [-1, 29 ]
+    }
+]
+
+instance_group {
+  count: 1#NUM_ENGINES
+  gpus: 0
+  kind: KIND_GPU
+}
+
+#db#dynamic_batching {
+#db#    preferred_batch_size: 8#MAX_BATCH
+#db#    max_queue_delay_microseconds: #MAX_QUEUE
+#db#}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-tensorrt-ensemble/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-tensorrt-ensemble/config.pbtxt
new file mode 100644
index 0000000000..7fb601e8ca
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-tensorrt-ensemble/config.pbtxt
@@ -0,0 +1,63 @@
+name: "jasper-tensorrt-ensemble"
+platform: "ensemble"
+max_batch_size: 8#MAX_BATCH
+input {
+    name: "AUDIO_SIGNAL"
+    data_type: TYPE_FP32
+    dims: -1#AUDIO_LENGTH
+}
+input {
+    name: "NUM_SAMPLES"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
+}
+output {
+    name: "TRANSCRIPT"
+    data_type: TYPE_INT32
+    dims: [-1]
+}
+ensemble_scheduling {
+    step {
+	model_name: "feature-extractor-ts-trace"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_SIGNAL"
+	}
+	input_map {
+	    key: "input__1"
+	    value: "NUM_SAMPLES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "AUDIO_FEATURES"
+	}    
+    }
+    step {
+	model_name: "jasper-tensorrt"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_FEATURES"
+
+	}
+	output_map {
+	    key: "output__0"
+	    value: "CHARACTER_PROBABILITIES"
+
+	}
+
+    }
+    step {
+	model_name: "decoder-ts-script"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "CHARACTER_PROBABILITIES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "TRANSCRIPT"
+	}
+    }
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-tensorrt/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-tensorrt/config.pbtxt
new file mode 100644
index 0000000000..032482f9d4
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-tensorrt/config.pbtxt
@@ -0,0 +1,58 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# Redistribution and use in source and binary forms, with or without
+# modification, are permitted provided that the following conditions
+# are met:
+#  * Redistributions of source code must retain the above copyright
+#    notice, this list of conditions and the following disclaimer.
+#  * Redistributions in binary form must reproduce the above copyright
+#    notice, this list of conditions and the following disclaimer in the
+#    documentation and/or other materials provided with the distribution.
+#  * Neither the name of NVIDIA CORPORATION nor the names of its
+#    contributors may be used to endorse or promote products derived
+#    from this software without specific prior written permission.
+#
+# THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS ``AS IS'' AND ANY
+# EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE
+# IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR
+# PURPOSE ARE DISCLAIMED.  IN NO EVENT SHALL THE COPYRIGHT OWNER OR
+# CONTRIBUTORS BE LIABLE FOR ANY DIRECT, INDIRECT, INCIDENTAL, SPECIAL,
+# EXEMPLARY, OR CONSEQUENTIAL DAMAGES (INCLUDING, BUT NOT LIMITED TO,
+# PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES; LOSS OF USE, DATA, OR
+# PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND ON ANY THEORY
+# OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+# (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE
+# OF THIS SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+
+name: "jasper-tensorrt"
+platform: "tensorrt_plan"
+default_model_filename: "model.plan"
+
+max_batch_size: 8#MAX_BATCH
+
+input [
+    {
+        name: "input__0"
+        data_type: TYPE_FP32
+        dims: [64, -1]
+    }      
+]
+
+output [
+    {
+        name: "output__0"
+        data_type: TYPE_FP32
+        dims: [-1, 29 ]
+    }
+]
+
+instance_group {
+  count: 1#NUM_ENGINES
+  gpus: 0
+  kind: KIND_GPU
+}
+
+#db#dynamic_batching {
+#db#    preferred_batch_size: 8#MAX_BATCH
+#db#    max_queue_delay_microseconds: #MAX_QUEUE
+#db#}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-ts-trace-ensemble/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-ts-trace-ensemble/config.pbtxt
new file mode 100644
index 0000000000..5da2ea9c0f
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-ts-trace-ensemble/config.pbtxt
@@ -0,0 +1,63 @@
+name: "jasper-ts-trace-ensemble"
+platform: "ensemble"
+max_batch_size: 8#MAX_BATCH
+input {
+    name: "AUDIO_SIGNAL"
+    data_type: TYPE_FP32
+    dims: -1#AUDIO_LENGTH
+}
+input {
+    name: "NUM_SAMPLES"
+    data_type: TYPE_INT32
+    dims: [ 1 ]
+}
+output {
+    name: "TRANSCRIPT"
+    data_type: TYPE_INT32
+    dims: [-1]
+}
+ensemble_scheduling {
+    step {
+	model_name: "feature-extractor-ts-trace"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_SIGNAL"
+	}
+	input_map {
+	    key: "input__1"
+	    value: "NUM_SAMPLES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "AUDIO_FEATURES"
+	}    
+    }
+    step {
+	model_name: "jasper-ts-trace"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "AUDIO_FEATURES"
+
+	}
+	output_map {
+	    key: "output__0"
+	    value: "CHARACTER_PROBABILITIES"
+
+	}
+
+    }
+    step {
+	model_name: "decoder-ts-script"
+	model_version: -1
+	input_map {
+	    key: "input__0"
+	    value: "CHARACTER_PROBABILITIES"
+	}
+	output_map {
+	    key: "output__0"
+	    value: "TRANSCRIPT"
+	}
+    }
+}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-ts-trace/config.pbtxt b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-ts-trace/config.pbtxt
new file mode 100644
index 0000000000..2bb708f606
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/model_repo_configs/fp32/jasper-ts-trace/config.pbtxt
@@ -0,0 +1,32 @@
+name: "jasper-ts-trace"
+platform: "pytorch_libtorch"
+default_model_filename: "model.pt"
+
+max_batch_size: 8#MAX_BATCH
+
+input [
+  {
+    name: "input__0"
+    data_type: TYPE_FP32
+    dims: [64, -1]
+  }
+]
+
+output [
+  {
+    name: "output__0"
+    data_type: TYPE_FP32
+    dims: [-1, 29]
+  }
+]
+
+instance_group {
+  count: 1#NUM_ENGINES
+  gpus: 0
+  kind: KIND_GPU
+}
+
+#db#dynamic_batching {
+#db#    preferred_batch_size: 8#MAX_BATCH
+#db#    max_queue_delay_microseconds: #MAX_QUEUE
+#db#}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/pytorch/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/pytorch/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/pytorch/utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/pytorch/utils.py
new file mode 100644
index 0000000000..823317b8b9
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/pytorch/utils.py
@@ -0,0 +1,263 @@
+# *****************************************************************************
+#  Copyright (c) 2018, NVIDIA CORPORATION.  All rights reserved.
+#
+#  Redistribution and use in source and binary forms, with or without
+#  modification, are permitted provided that the following conditions are met:
+#      * Redistributions of source code must retain the above copyright
+#        notice, this list of conditions and the following disclaimer.
+#      * Redistributions in binary form must reproduce the above copyright
+#        notice, this list of conditions and the following disclaimer in the
+#        documentation and/or other materials provided with the distribution.
+#      * Neither the name of the NVIDIA CORPORATION nor the
+#        names of its contributors may be used to endorse or promote products
+#        derived from this software without specific prior written permission.
+#
+#  THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
+#  ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
+#  WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
+#  DISCLAIMED. IN NO EVENT SHALL NVIDIA CORPORATION BE LIABLE FOR ANY
+#  DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
+#  (INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
+#  LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
+#  ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
+#  (INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
+#  SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
+#
+# *****************************************************************************
+
+import tensorrt as trt
+import torch
+from collections import Counter
+import json
+
+import logging
+
+triton_type_to_torch_type = {
+    'TYPE_BOOL': torch.bool,
+    'TYPE_INT8': torch.int8,
+    'TYPE_INT16': torch.int16,
+    'TYPE_INT32': torch.int32,
+    'TYPE_INT64': torch.int64,
+    'TYPE_UINT8': torch.uint8,
+    'TYPE_FP16': torch.float16,
+    'TYPE_FP32': torch.float32,
+    'TYPE_FP64': torch.float64
+}
+
+torch_type_to_triton_type = {
+    torch.bool: 'TYPE_BOOL',
+    torch.int8: 'TYPE_INT8',
+    torch.int16: 'TYPE_INT16',
+    torch.int32: 'TYPE_INT32',
+    torch.int64: 'TYPE_INT64',
+    torch.uint8: 'TYPE_UINT8',
+    torch.float16: 'TYPE_FP16',
+    torch.float32: 'TYPE_FP32',
+    torch.float64: 'TYPE_FP64'
+}
+
+
+def build_tensorrt_engine(model_file, shapes, max_workspace_size,
+                          max_batch_size, fp16_mode):
+    ''' takes a path to an onnx file, and shape information, returns a tensorrt engine
+        :: model_file :: path to an onnx model
+        :: shapes :: dictionary containing min shape, max shape, opt shape for the tensorrt engine
+    '''
+    TRT_LOGGER = trt.Logger(trt.Logger.WARNING)
+    builder = trt.Builder(TRT_LOGGER)
+    builder.fp16_mode = fp16_mode
+    builder.max_batch_size = max_batch_size
+    #
+    config = builder.create_builder_config()
+    config.max_workspace_size = max_workspace_size
+
+    if fp16_mode:
+        config.flags |= 1 << int(trt.BuilderFlag.FP16)
+    profile = builder.create_optimization_profile()
+    for s in shapes:
+        profile.set_shape(s['name'], min=s['min'], opt=s['opt'], max=s['max'])
+    config.add_optimization_profile(profile)
+    explicit_batch = 1 << (int)(trt.NetworkDefinitionCreationFlag.EXPLICIT_BATCH)
+    network = builder.create_network(explicit_batch)
+    #
+    with trt.OnnxParser(network, TRT_LOGGER) as parser:
+        with open(model_file, 'rb') as model:
+            parser.parse(model.read())
+            for i in range(parser.num_errors):
+                print("[Converter error]: OnnxParser:", parser.get_error(i))
+            engine = builder.build_engine(network, config=config)
+    return engine
+
+
+def get_inputs(dataloader, device, precision):
+    ''' load sample inputs to device '''
+    inputs = []
+    logging.info("Loading sample inputs to device.")
+    for idx, batch in enumerate(dataloader):
+        if idx % (len(dataloader)//100) == 0:
+            logging.info(f"{idx}/{len(dataloader)}")
+
+        if type(batch) is torch.Tensor:
+            batch_d = batch.to(device)
+            if batch_d.is_floating_point() and precision == 'fp16':
+                batch_d = batch_d.to(torch.float16)
+            batch_d = (batch_d,)
+            inputs.append(batch_d)
+        else:
+            batch_d = []
+            for x in batch:
+                assert type(x) is torch.Tensor, "input is not a tensor"
+                x = x.to(device)
+                if x.is_floating_point() and precision == 'fp16':
+                    x = x.to(torch.float16)
+                batch_d.append(x)
+            batch_d = tuple(batch_d)
+            inputs.append(batch_d)
+    logging.info("Finished loading sample inputs to device.")
+    return inputs
+
+
+def get_list_of_shapes(l, fun):
+    ''' returns the list of min/max shapes, depending on fun
+        :: l :: list of tuples of tensors
+        :: fun :: min or max
+    '''
+    tensor_tuple = l[0]
+    shapes = [list(x.shape) for x in tensor_tuple]
+    for tensor_tuple in l:
+        assert len(tensor_tuple) == len(shapes), "tensors with varying shape lengths are not supported"
+        for i,x in enumerate(tensor_tuple):
+            for j in range(len(x.shape)):
+                shapes[i][j] = fun(shapes[i][j], x.shape[j])
+    return shapes # a list of shapes
+
+
+def get_min_shapes(l):
+    ''' returns the tuple of min shapes
+        :: l :: list of tuples of tensors '''
+    shapes = get_list_of_shapes(l, min)
+    min_batch = 1
+    shapes = [[min_batch,*shape[1:]] for shape in shapes]
+    shapes = tuple(shapes)
+    return shapes # tuple of min shapes
+
+
+def get_max_shapes(l):
+    ''' returns the tuple of max shapes
+        :: l :: list of tuples of tensors '''
+    shapes = get_list_of_shapes(l, max)
+    max_batch = max(1,shapes[0][0])
+    shapes = [[max_batch,*shape[1:]] for shape in shapes]
+    shapes = tuple(shapes)
+    return shapes # tuple of max shapes
+
+
+def get_opt_shapes(l):
+    ''' returns the tuple of opt shapes
+        :: l :: list of tuples of tensors '''
+    counter = Counter()
+    for tensor_tuple in l:
+        shapes = [tuple(x.shape) for x in tensor_tuple]
+        shapes = tuple(shapes)
+        counter[shapes] += 1
+    shapes = counter.most_common(1)[0][0]
+    return shapes # tuple of most common occuring shapes
+
+
+def get_shapes(l, max_batch_size):
+    ''' returns a tuple of dynamic shapes: variable tensor dimensions
+        (for ex. batch size) occur as -1 in the tuple
+        :: l :: list of tuples of tensors '''
+    tensor_tuple = l[0]
+    shapes = [list(x.shape) for x in tensor_tuple]
+    for tensor_tuple in l:
+        err_msg = "tensors with varying shape lengths are not supported"
+        assert len(tensor_tuple) == len(shapes), err_msg
+        for i,x in enumerate(tensor_tuple):
+            for j in range(len(x.shape)):
+                if shapes[i][j] != x.shape[j] or j == 0 and max_batch_size > 1:
+                    shapes[i][j] = -1
+    shapes = tuple(shapes)
+    return shapes # tuple of dynamic shapes
+
+
+def get_io_properties(inputs, outputs, max_batch_size):
+
+    # generate input shapes - dynamic tensor shape support
+    input_shapes = get_shapes(inputs, max_batch_size)
+
+    # generate output shapes - dynamic tensor shape support
+    output_shapes = get_shapes(outputs, max_batch_size)
+
+    # generate input types
+    input_types = [torch_type_to_triton_type[x.dtype] for x in inputs[0]]
+
+    # generate output types
+    output_types = [torch_type_to_triton_type[x.dtype] for x in outputs[0]]
+
+    # get input names
+    rng = range(len(input_types))
+    input_names = ["input__" + str(num) for num in rng]
+
+    # get output names
+    rng = range(len(output_types))
+    output_names = ["output__" + str(num) for num in rng]
+
+    # get indices of dynamic input and output shapes
+    dynamic_axes = {}
+    for input_name,input_shape in zip(input_names,input_shapes):
+        dynamic_axes[input_name] = [i for i,x in enumerate(input_shape) if x == -1]
+    for output_name,output_shape in zip(output_names,output_shapes):
+        dynamic_axes[output_name] = [i for i,x in enumerate(output_shape) if x == -1]
+
+    # min, opt, max shapes for TensorRT
+    min_shapes = get_min_shapes(inputs)
+    opt_shapes = get_opt_shapes(inputs)
+    max_shapes = get_max_shapes(inputs)
+
+    res = {"input_shapes": input_shapes,
+            "output_shapes": output_shapes,
+            "input_types": input_types,
+            "output_types": output_types,
+            "input_names": input_names,
+            "output_names": output_names,
+            "dynamic_axes": dynamic_axes,
+            "min_shapes": min_shapes,
+            "opt_shapes": opt_shapes,
+            "max_shapes": max_shapes}
+
+    return res
+
+
+def extract_io_props(model, dataloader, device, precision, max_batch_size):
+
+    # prepare inputs
+    inputs = get_inputs(dataloader, device, precision)
+    # generate outputs
+    outputs = []
+    for input in inputs:
+        with torch.no_grad():
+            output = model(*input)
+        if type(output) is torch.Tensor:
+            output = [output]
+        outputs.append(output)
+
+    # prepare input/output properties
+    io_props = get_io_properties(inputs, outputs, max_batch_size)
+
+    return io_props
+
+def save_io_props(io_props, io_props_path):
+
+    with open(io_props_path, "w") as f:
+        f.write(json.dumps(io_props))
+
+
+def load_io_props(io_props_path):
+
+    with open(io_props_path, "r") as f:
+        data = json.loads(f.read())
+    if "dynamic_axes" not in data.keys():
+        return data
+
+    return data
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/docker/build_triton_client.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/docker/build_triton_client.sh
new file mode 100644
index 0000000000..6114cb81d8
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/docker/build_triton_client.sh
@@ -0,0 +1,7 @@
+#!/bin/bash
+# ensure the TRTIS submodule is added and build the clients
+SCRIPT_DIR=$(cd $(dirname $0); pwd)
+PROJECT_DIR=${SCRIPT_DIR}/../../../
+docker pull nvcr.io/nvidia/tritonserver:20.10-py3-clientsdk
+git submodule update --init --recursive
+docker build . --rm -f ${PROJECT_DIR}/triton/Dockerfile -t jasper:triton
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/download_triton_librispeech.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/download_triton_librispeech.sh
new file mode 100644
index 0000000000..c6941e8caa
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/download_triton_librispeech.sh
@@ -0,0 +1,29 @@
+#!/usr/bin/env bash
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+#Downloads the inference-subset of the Librispeech corpus.
+
+
+DATA_SET="LibriSpeech"
+DATA_ROOT_DIR="/datasets"
+DATA_DIR="${DATA_ROOT_DIR}/${DATA_SET}"
+if [ ! -d "$DATA_DIR" ]
+then
+    mkdir -p $DATA_DIR
+    chmod go+rx $DATA_DIR
+    python utils/download_librispeech.py triton/triton_librispeech.csv $DATA_DIR -e ${DATA_ROOT_DIR}/
+else
+    echo "Directory $DATA_DIR already exists."
+fi
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/execute_all_perf_runs.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/execute_all_perf_runs.sh
new file mode 100644
index 0000000000..892c011bc0
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/execute_all_perf_runs.sh
@@ -0,0 +1,77 @@
+#!/bin/bash
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+# This will run all the necessary scripts to generate ALL needed output
+
+if [ -f /.dockerenv ]; then # inside docker
+    echo "The script \"$0\" script should be run from outside the container. Exiting."
+    exit 1
+fi
+
+#### input arguments
+RESULT_DIR=${RESULT_DIR}
+####
+
+for dir in $CHECKPOINT_DIR $RESULT_DIR; do
+    if [[ $dir != /* ]]; then
+        echo "All directory paths must be absolute paths!"
+        echo "${dir} is not an absolute path"
+        exit 1
+    fi
+
+    if [ ! -d $dir ]; then
+        echo "All directory paths must exist!"
+        echo "${dir} does not exist"
+        exit 1
+    fi
+done
+
+export GPU=${GPU:-}
+
+SCRIPT_DIR=$(cd $(dirname $0); pwd)
+PROJECT_DIR=${SCRIPT_DIR}/../..
+MODEL_REPO=${MODEL_REPO:-"${PROJECT_DIR}/triton/deploy/model_repo"}
+
+# We need to make sure TRITON uses only one GPU, same as export does
+# for TRITON
+export NVIDIA_VISIBLE_DEVICES=${GPU}
+export TRITON_CLIENT_CONTAINER_TAG=jasper:triton
+
+trap "exit" INT
+
+
+function run_for_length () {
+    TRITON_CLIENT_CONTAINER_TAG=jasper:triton \
+			       AUDIO_LENGTH=$1 \
+			       BATCH_SIZE=$2 \
+			       RESULT_DIR=${RESULT_DIR} \
+			       PRECISION=${PRECISION} \
+			       ${SCRIPT_DIR}/generate_perf_results.sh
+}
+
+PRECISION_TESTS=${PRECISION_TESTS:-"fp16 fp32"}
+BATCH_SIZES=${BATCH_SIZES:-"1" "2" "4" "8"}
+SEQ_LENS=${SEQ_LENS:-"32000" "112000" "267200"} # i.e. ,7, 2, and 16.7 seconds
+
+for PRECISION in ${PRECISION_TESTS};
+do
+    for BATCH_SIZE in ${BATCH_SIZES};
+    do
+	for LENGTH in ${SEQ_LENS};
+	do
+            run_for_length $LENGTH $BATCH_SIZE
+	done
+    done
+done
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/export_model.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/export_model.sh
new file mode 100644
index 0000000000..4db732cafb
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/export_model.sh
@@ -0,0 +1,78 @@
+#!/bin/bash
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+SCRIPT_DIR=$(cd $(dirname $0); pwd)
+PROJECT_DIR=${SCRIPT_DIR}/../..
+TRITON_DIR=${PROJECT_DIR}/triton
+export PYTHONPATH=$TRITON_DIR:$PYTHONPATH
+
+# Jasper model arguments ---------------------------------------------
+BATCH_SIZE=${BATCH_SIZE:-8}
+DATA_DIR=${DATA_DIR:-"/datasets/LibriSpeech"}
+DATASET=${DATASET:-"test-clean"}
+MODEL_CONFIG=${MODEL_CONFIG:-"configs/jasper10x5dr_speedp-offline_speca_nomask.yaml"}
+CHECKPOINT_PATH=${CHECKPOINT_PATH:-"/checkpoints/nvidia_jasper_210205.pt"}
+EMA=${EMA:-true}
+
+ARGS=" --batch_size $BATCH_SIZE "
+ARGS+=" --dataset_dir $DATA_DIR "
+ARGS+=" --val_manifest $DATA_DIR/librispeech-${DATASET}-wav.json "
+ARGS+=" --model_config $MODEL_CONFIG  "
+[ -n "$CHECKPOINT_PATH" ] && \
+    ARGS+=" --ckpt=${CHECKPOINT_PATH}"
+[ "$EMA" = true ] &&                 ARGS+=" --ema"
+## -------------------------------------------------------------------
+
+run_converter() {
+    set -x
+    mkdir -p $3
+    chmod a+w $3
+    python triton/converter.py --model-module jasper_module \
+    	   --convert $1 \
+    	   --precision ${PRECISION} \
+    	   --save-dir $3 \
+    	   -- \
+    	   ${ARGS} --component $2
+    set +x
+}
+
+CONVERTS=${CONVERTS:-"feature-extractor" "decoder" "ts-trace" "onnx" "tensorrt"}
+CONVERT_PRECISIONS=${CONVERT_PRECISIONS:-"fp16" "fp32"}
+
+for PRECISION in ${CONVERT_PRECISIONS[@]}; do
+
+    MODEL_REPO_DIR="${TRITON_DIR}/model_repo/${PRECISION}/"
+    mkdir -p $MODEL_REPO_DIR
+    chmod a+w "${TRITON_DIR}" "${TRITON_DIR}/model_repo/" "${TRITON_DIR}/model_repo/${PRECISION}/"
+
+    if [[ " ${CONVERTS[@]} " =~ " feature-extractor " ]]; then
+    	## Export data preprocessor
+    	run_converter "ts-trace" "feature-extractor" \
+		      ${MODEL_REPO_DIR}/feature-extractor-ts-trace/1
+    fi
+    if [[ " ${CONVERTS[@]} " =~ " decoder " ]]; then
+    	## Export greedy decoder
+    	run_converter "ts-script" "decoder" \
+		      ${MODEL_REPO_DIR}/decoder-ts-script/1
+    fi
+    BACKENDS=( "ts-trace" "onnx" "tensorrt" )
+    for BACKEND in ${BACKENDS[@]}; do
+	if [[ " ${CONVERTS[@]} " =~ " ${BACKEND} " ]]; then
+            ## Export acoustic model - Jasper
+            run_converter "${BACKEND}" "model" \
+			  ${MODEL_REPO_DIR}/jasper-${BACKEND}/1
+	fi
+    done
+done
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/generate_perf_results.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/generate_perf_results.sh
new file mode 100644
index 0000000000..7d57834c75
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/generate_perf_results.sh
@@ -0,0 +1,108 @@
+#!/bin/bash
+
+#### input arguments
+RESULT_DIR=${RESULT_DIR} # used by perf_client to store results
+NVIDIA_VISIBLE_DEVICES=${NVIDIA_VISIBLE_DEVICES:-"0"}
+SERVER_HOSTNAME=${SERVER_HOSTNAME:-localhost}
+AUDIO_LENGTH=${AUDIO_LENGTH:-80000}
+BATCH_SIZE=${BATCH_SIZE:-16}
+####
+
+set -e
+SCRIPT_DIR=$(cd $(dirname $0); pwd)
+TRITON_DIR=${SCRIPT_DIR}/..
+PROJECT_DIR=${TRITON_DIR}/..
+
+GPU_TESTS=${GPU_TESTS:-"tensorrt ts-trace onnx"}
+ENGINE_COUNT_TESTS=${ENGINE_COUNT_TESTS:-"1"}
+
+# Export the set variables in case they used the default
+export NVIDIA_VISIBLE_DEVICES
+export SERVER_HOSTNAME
+export AUDIO_LENGTH
+export MAX_LATENCY=2000 # Set max latency high to prevent errors
+TRITON=${TRITON:-jasper-triton-server}
+MAX_QUEUE_DELAYS=${MAX_QUEUE_DELAYS:-"10 5 2"} #ms
+
+# Ensure that the server is closed when the script exits
+function cleanup_server {
+    current_time=$(date "+%Y.%m.%d-%H.%M.%S")
+    logfile="/tmp/${TRITON}-${current_time}.log"
+    echo "Shutting down ${TRITON} container, log is in ${logfile}"
+    docker logs ${TRITON} > ${logfile} 2>&1
+    docker stop ${TRITON} > /dev/null 2>&1
+}
+
+trap cleanup_server EXIT
+
+trap "exit" INT
+
+function wait_for_triton {
+    TIMEOUT=${1:-60}
+    timeout ${TIMEOUT} ${SCRIPT_DIR}/wait_for_triton_server.sh || (echo '\nServer timeout!!!\n' && exit 1)
+}
+
+function modify_ensemble {
+
+    PLAT=$1
+
+    REPO=${TRITON_DIR}/deploy/model_repo
+    INPLACE="--in_place"
+
+    CONF=${REPO}/jasper-${PLAT}/config.pbtxt
+    CONF_E=${REPO}/jasper-${PLAT}-ensemble/config.pbtxt
+
+
+    echo "Modifying ${CONF} : batch size ${BATCH_SIZE} engines=${NUM_ENGINES} ..."
+    cleanup_server || true
+    sed -i -e "s/1#NUM_ENGINES/${NUM_ENGINES}/g" -e "s/8#MAX_BATCH/${BATCH_SIZE}/g" ${CONF}
+    if [ "$MAX_QUEUE" != "" ] ; then
+        sed -i -e "s/#db#//g" -e "s/#MAX_QUEUE/${MAX_QUEUE}/g" ${CONF}
+    fi
+
+
+    echo "Modifying ${CONF_E} for size $2, batch size ${BATCH_SIZE} ${TRITON_DYN_BATCH_ARGS}.."
+    sed -i -e "s/-1#AUDIO_LENGTH/${AUDIO_LENGTH}/g" -e "s/8#MAX_BATCH/${BATCH_SIZE}/g" ${CONF_E}
+
+    ${SCRIPT_DIR}/run_server.sh
+
+    wait_for_triton
+
+    echo "done."
+}
+
+echo "GPU tests: ${GPU_TESTS}"
+echo "PRECISION: ${PRECISION}"
+
+for plat in ${GPU_TESTS}; do
+    if [ "$plat" == "none" ]; then
+	    continue
+    else
+        export MAX_LATENCY=2000
+        export MEASUREMENT_WINDOW=3000
+    fi
+
+    export BASE_SAVE_NAME="${plat}_${PRECISION}_${AUDIO_LENGTH}_BS${BATCH_SIZE}"
+    export MODEL_NAME=jasper-${plat}-ensemble
+
+    MODELS="jasper-${plat} jasper-${plat}-ensemble" PRECISION=${PRECISION} ${SCRIPT_DIR}/prepare_model_repository.sh
+
+    ############## Engine Count Comparison (static batcing) ##############
+    for num_engines in ${ENGINE_COUNT_TESTS}; do
+	SAVE_RESULTS_DIR="${BASE_SAVE_NAME}/static/${num_engines}_engines"
+	NUM_ENGINES=${num_engines} BATCH_SIZE=${BATCH_SIZE} modify_ensemble ${plat} ${AUDIO_LENGTH}
+	echo "Running engines comparison, ${num_engines} engines..."
+	MAX_CONCURRENCY=8 BATCH_SIZE=${BATCH_SIZE} ${SCRIPT_DIR}/run_perf_client.sh  ${SAVE_RESULTS_DIR} || echo '\nPerf Client Failure!!!\n'
+    done
+
+    ############## Dynamic Batching Comparison ##############
+    for delay in ${MAX_QUEUE_DELAYS}; do
+	echo "Running dynamic batching comparison, models=${MODELS}, delay ${delay}..."
+	TRITON_DYN_BATCHING_DELAY=$((delay * 1000))
+	SAVE_RESULTS_DIR="${BASE_SAVE_NAME}/batching/${TRITON_DYN_BATCHING_DELAY}"
+	NUM_ENGINES=1 MAX_QUEUE=${TRITON_DYN_BATCHING_DELAY} BATCH_SIZE=${BATCH_SIZE} modify_ensemble ${plat} ${AUDIO_LENGTH}
+	BATCH_SIZE=1 MAX_CONCURRENCY=$((BATCH_SIZE*2)) ${SCRIPT_DIR}/run_perf_client.sh ${SAVE_RESULTS_DIR} || echo '\nPerf Client Failure!!!\n'
+    done
+done
+
+echo "Complete!"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/prepare_model_repository.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/prepare_model_repository.sh
new file mode 100644
index 0000000000..6608718c23
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/prepare_model_repository.sh
@@ -0,0 +1,54 @@
+#!/bin/bash
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+
+# Create folder deploy/model_repo that will be used by TRITON
+
+SCRIPT_DIR=$(cd $(dirname $0); pwd)
+PROJECT_DIR=${SCRIPT_DIR}/..
+DEPLOY_DIR=${PROJECT_DIR}/deploy
+HOST_REPO=${DEPLOY_DIR}/model_repo
+
+
+MODELS_TENSORRT=${MODELS_TENSORRT:-"jasper-tensorrt jasper-tensorrt-ensemble"}
+MODELS_TS_TRACE=${MODELS_TS_TRACE:-"jasper-ts-trace jasper-ts-trace-ensemble"}
+MODELS_ONNX=${MODELS_ONNX:-"jasper-onnx jasper-onnx-ensemble"}
+DECODERS="decoder-ts-script"
+EXTRACTORS="feature-extractor-ts-trace"
+
+MODELS=${MODELS:-"${MODELS_ONNX} ${MODELS_TENSORRT} ${MODELS_TS_TRACE}"}
+PRECISION=${PRECISION:-"fp16" "fp32"}
+
+# only link working models to install directory
+rm -fr ${HOST_REPO} && mkdir -p ${HOST_REPO}
+if [ -f /.dockerenv ]; then # inside docker
+    chmod -R a+w ${HOST_REPO}
+fi
+
+echo "Setting up model repo at ${HOST_REPO}, models: ${MODELS} ..."
+for m  in ${EXTRACTORS} ${DECODERS} ${MODELS}; do
+
+    mkdir -p ${HOST_REPO}/$m
+    cp ${PROJECT_DIR}/model_repo_configs/${PRECISION}/$m/config.pbtxt ${HOST_REPO}/$m/
+    if [ -d "${PROJECT_DIR}/model_repo/${PRECISION}/$m/1"  ]; then
+	echo "Creating symlink ls -sf /model_repo/${PRECISION}/$m/1 ${HOST_REPO}/$m"
+	ln -sf /model_repo/${PRECISION}/$m/1 ${HOST_REPO}/$m
+    else
+	mkdir -p ${HOST_REPO}/$m/1
+    fi
+    if [ -f /.dockerenv ]; then # inside docker
+	chmod -R a+w ${HOST_REPO}/$m
+    fi
+done
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/preprocess_triton_librispeech.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/preprocess_triton_librispeech.sh
new file mode 100644
index 0000000000..e62757dd76
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/preprocess_triton_librispeech.sh
@@ -0,0 +1,21 @@
+#!/usr/bin/env bash
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+# Constructs JSON manifests for inference-subset of Librispeech corpus.
+
+python ./utils/convert_librispeech.py \
+    --input_dir /datasets/LibriSpeech/test-clean \
+    --dest_dir /datasets/LibriSpeech/test-clean-wav \
+    --output_json /datasets/LibriSpeech/librispeech-test-clean-wav.json
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_client.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_client.sh
new file mode 100644
index 0000000000..129acbd4fe
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_client.sh
@@ -0,0 +1,49 @@
+#!/bin/bash
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+SCRIPT_DIR=$(cd $(dirname $0); pwd)
+PROJECT_DIR=${SCRIPT_DIR}/../..
+
+MODEL_TYPE=${1:-"ts-trace"}
+DATA_DIR=${2} # folder with data
+FILE=${3} # json manifest file, OR single wav file
+
+JASPER_CONTAINER_TAG=${JASPER_CONTAINER_TAG:-jasper:triton}
+
+if [ "$#" -ge 1 ] && [ "${FILE: -4}" == ".wav" ]; then 
+  CMD="python /jasper/triton/jasper-client.py --data_dir /data --audio_filename ${FILE} --model_platform ${MODEL_TYPE}"
+  ARGS=""
+  ARGS="$ARGS -v $DATA_DIR:/data"
+elif [ "$#" -ge 1 ] && [ "${FILE: -4}" == "json" ]; then
+  ARGS=""
+  ARGS="$ARGS -v $DATA_DIR:/data"
+  CMD="python /jasper/triton/jasper-client.py --manifest_filename ${FILE} --model_platform ${MODEL_TYPE} --data_dir /data"
+else
+  ARGS="-it"
+  CMD=""
+fi
+
+echo "========== STARTING ${JASPER_CONTAINER_TAG} =========="
+
+set -x
+nvidia-docker run --rm -it \
+   --net=host \
+   --shm-size=1g \
+   --ulimit memlock=-1 \
+   --ulimit stack=67108864 \
+   -v ${PROJECT_DIR}:/jasper \
+   --name=jasper-triton-client \
+   ${ARGS} ${JASPER_CONTAINER_TAG} ${CMD}
+set +x
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_perf_client.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_perf_client.sh
new file mode 100644
index 0000000000..43c1382658
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_perf_client.sh
@@ -0,0 +1,86 @@
+
+#!/bin/bash
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+trap "exit" INT
+
+SCRIPT_DIR=$(cd $(dirname $0); pwd)
+PROJECT_DIR=${SCRIPT_DIR}/../..
+
+TRITON_CLIENT_CONTAINER_TAG=${TRITON_CLIENT_CONTAINER_TAG:-jasper:triton}
+
+SERVER_HOSTNAME=${SERVER_HOSTNAME:-localhost}
+MODEL_NAME=${MODEL_NAME:-jasper-tensorrt-ensemble}
+MODEL_VERSION=${MODEL_VERSION:-1}
+BATCH_SIZE=${BATCH_SIZE:-1}
+AUDIO_LENGTH=${AUDIO_LENGTH:-32000}
+RESULT_DIR=${RESULT_DIR:-${PROJECT_DIR}/results}
+MAX_LATENCY=${MAX_LATENCY:-500}
+MAX_CONCURRENCY=${MAX_CONCURRENCY:-64}
+MEASUREMENT_WINDOW=${MEASUREMENT_WINDOW:-3000}
+
+TIMESTAMP=$(date "+%y%m%d_%H%M")
+
+# RESULT_DIR_H is the path on the host, outside the container. Inside the container RESULT_DIR_H is always mounted to /results
+RESULT_DIR_H="${RESULT_DIR}/perf_client/${MODEL_NAME}"
+
+# Set the output folder using the first argument or pick a default
+if [ -z ${1+x} ]; then
+   RESULT_DIR_H=${RESULT_DIR_H}/batch_${BATCH_SIZE}_len_${AUDIO_LENGTH}
+else
+   RESULT_DIR_H=${RESULT_DIR_H}/"$1"
+   shift
+fi
+
+# Make the directory if it doesnt exist
+if [ ! -d "${RESULT_DIR_H}" ]; then
+   mkdir -p "${RESULT_DIR_H}"
+fi
+
+echo "Saving output to ${RESULT_DIR_H}"
+
+LOGNAME="${RESULT_DIR_H}/log_${TIMESTAMP}.log"
+OUTPUT_FILE_CSV="results_${TIMESTAMP}.csv"
+
+ARGS="\
+   -m ${MODEL_NAME} \
+   -x ${MODEL_VERSION} \
+   -p ${MEASUREMENT_WINDOW} \
+   -v \
+   -i gRPC \
+   -u ${SERVER_HOSTNAME}:8001 \
+   -b ${BATCH_SIZE} \
+   -l ${MAX_LATENCY} \
+   --max-threads ${MAX_CONCURRENCY} "
+
+curl -s "http://${SERVER_HOSTNAME}:8000/api/status/${MODEL_NAME}" | grep ready_state | grep SERVER_READY || (echo "Model ${MODEL_NAME} is not ready, perf_client skipped..." && exit 1)
+
+echo "=== STARTING: perf client ${ARGS} --concurrency-range 1:4:1 ==="
+set -x
+docker run  -e DISPLAY=${DISPLAY}  --runtime nvidia --rm \
+	      --privileged --net=host \
+	      -v ${RESULT_DIR_H}:/results --name jasper-perf-client \
+	      ${TRITON_CLIENT_CONTAINER_TAG}  perf_client $ARGS -f /results/${OUTPUT_FILE_CSV}_p1 --shape AUDIO_SIGNAL:${AUDIO_LENGTH} --concurrency-range 1:4:1 2>&1 | tee -a $LOGNAME
+set +x
+
+echo "=== STARTING: perf client ${ARGS} --concurrency-range 8:${MAX_CONCURRENCY}:8 ==="
+set -x
+docker run  -e DISPLAY=${DISPLAY}  --runtime nvidia --rm \
+	      --privileged --net=host \
+	      -v ${RESULT_DIR_H}:/results --name jasper-perf-client \
+	      ${TRITON_CLIENT_CONTAINER_TAG}  perf_client $ARGS -f /results/${OUTPUT_FILE_CSV}_p2 --shape AUDIO_SIGNAL:${AUDIO_LENGTH} --concurrency-range 8:${MAX_CONCURRENCY}:8 2>&1 | tee -a $LOGNAME
+set +x
+
+cat ${RESULT_DIR_H}/${OUTPUT_FILE_CSV}_p1 ${RESULT_DIR_H}/${OUTPUT_FILE_CSV}_p2 > ${RESULT_DIR_H}/${OUTPUT_FILE_CSV}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_server.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_server.sh
new file mode 100644
index 0000000000..03360c9962
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/run_server.sh
@@ -0,0 +1,72 @@
+#!/bin/bash
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+SCRIPT_DIR=$(cd $(dirname $0); pwd)
+TRITON_DIR=${SCRIPT_DIR}/..
+
+DEPLOY_DIR=${DEPLOY_DIR:-${TRITON_DIR}/deploy}
+MODELS_DIR=${MODEL_DIR:-"$DEPLOY_DIR/model_repo"}
+TRITON_CONTAINER_TAG=${TRITON_CONTAINER_TAG:-"nvcr.io/nvidia/tritonserver:20.10-py3"}
+NV_VISIBLE_DEVICES=${NVIDIA_VISIBLE_DEVICES:-"0"} # FIXME "all"
+TRITON=${TRITON:-jasper-triton-server}
+
+MODEL_TYPE=${1:-""}
+PRECISION=${2:-""}
+
+# Ensure that the server is closed when the script exits
+function cleanup_server {
+    current_time=$(date "+%Y.%m.%d-%H.%M.%S")
+    logfile="/tmp/${TRITON}-${current_time}.log"
+    echo "Shutting down ${TRITON} container, log is in ${logfile}"
+    docker logs ${TRITON} > ${logfile} 2>&1
+    docker stop ${TRITON} > /dev/null 2>&1
+}
+
+trap "exit" INT
+
+# FIXME
+# if [ "$(docker inspect -f "{{.State.Running}}" ${TRITON})" = "true" ]; then
+#     if [ "$1" == "norestart" ]; then
+#        echo "${TRITON} is already running ..."
+#        exit 0
+#     fi
+#     cleanup_server || true
+# fi
+
+if [ -n "$MODEL_TYPE" ] && [ -n "$PRECISION" ]; then
+    MODELS="jasper-${MODEL_TYPE} jasper-${MODEL_TYPE}-ensemble" PRECISION=${PRECISION} ${SCRIPT_DIR}/prepare_model_repository.sh
+fi
+
+# To start  TRITON container with alternative commandline, set CMD
+CMD=${CMD:-"/opt/tritonserver/bin/tritonserver --log-verbose=100 --exit-on-error=true --strict-model-config=false --model-store=/models"}
+DAEMON=${DAEMON:-"-d"}
+RM=${RM:-"--rm"}
+
+if [ ! -d "${MODELS_DIR}" ]; then
+    echo "${MODELS_DIR} does not exist!"
+    echo "Exiting from script ${0}."
+    exit 1
+fi
+
+set -x
+docker run -p 8000:8000 -p 8001:8001 -p 8002:8002 \
+       --runtime nvidia \
+       -e NVIDIA_VISIBLE_DEVICES=${NV_VISIBLE_DEVICES} \
+       -v ${MODELS_DIR}:/models \
+       -v ${TRITON_DIR}/model_repo:/model_repo \
+       --shm-size=1g --ulimit memlock=-1 --ulimit stack=67108864  \
+       ${DAEMON} ${RM} --name ${TRITON} ${TRITON_CONTAINER_TAG} \
+       ${CMD}
+set +x
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/wait_for_triton_server.sh b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/wait_for_triton_server.sh
new file mode 100644
index 0000000000..1fb8a32afd
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/scripts/wait_for_triton_server.sh
@@ -0,0 +1,36 @@
+#!/bin/bash
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+SERVER_URI=${1:-"localhost"}
+
+echo "Waiting for TRITON Server to be ready at http://$SERVER_URI:8000..."
+
+
+live_command="curl -m 1 -L -s -o /dev/null -w %{http_code} localhost:8000/v2/health/live"
+ready_command="curl -m 1 -L -s -o /dev/null -w %{http_code} localhost:8000/v2/health/ready"
+
+current_status=$($live_command)
+
+SUCCESS_CODE=200
+
+# First check the current status. If that passes, check the json. If either fail, loop
+while [[ ${current_status} != ${SUCCESS_CODE} ]] || [[ $($ready_command) != ${SUCCESS_CODE} ]]; do
+
+   printf "."
+   sleep 1
+   current_status=$($live_command)
+done
+
+echo "TRITON Server is ready!"
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/speech_utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/speech_utils.py
new file mode 100644
index 0000000000..7f68c5b1ee
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/speech_utils.py
@@ -0,0 +1,417 @@
+#!/usr/bin/python
+
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+
+import soundfile as sf
+import math
+from os import system
+import numpy as np
+import tritonclient.grpc, tritonclient.http
+import tritonclient.grpc.model_config_pb2 as model_config
+from tritonclient.utils import triton_to_np_dtype, np_to_triton_dtype
+import grpc
+import sys
+import os
+if "./triton" not in sys.path:
+    sys.path.append(os.path.join(sys.path[0], "../"))
+from common.text import _clean_text
+
+WINDOWS_FNS = {"hanning": np.hanning, "hamming": np.hamming, "none": None}
+
+triton_type_to_np_dtype = {
+    'TYPE_BOOL': np.bool,
+    'TYPE_INT8': np.int8,
+    'TYPE_INT16': np.int16,
+    'TYPE_INT32': np.int32,
+    'TYPE_INT64': np.int64,
+    'TYPE_UINT8': np.uint8,
+    'TYPE_FP16': np.float16,
+    'TYPE_FP32': np.float32,
+    'TYPE_FP64': np.float64
+}
+
+model_dtype_to_np_dtype = {
+    "BOOL": np.bool,
+    "INT8": np.int8,
+    "INT16": np.int16,
+    "INT32": np.int32,
+    "INT64": np.int64,
+    "UINT8": np.uint8,
+    "UINT16": np.uint16,
+    "FP16": np.float16,
+    "FP32": np.float32,
+    "FP64": np.float64,
+    "BYTES": np.dtype(object)
+}
+
+def load_transcript(transcript_path):
+    with open(transcript_path, 'r', encoding="utf-8") as transcript_file:
+        transcript = transcript_file.read().replace('\n', '')
+    return transcript
+
+def parse_transcript(transcript, labels_map, blank_index):
+    chars = [labels_map.get(x, blank_index) for x in list(transcript)]
+    transcript = list(filter(lambda x: x != blank_index, chars))
+    return transcript
+
+def normalize_string(s, labels, table, **unused_kwargs):
+    """
+    Normalizes string. For example:
+    'call me at 8:00 pm!' -> 'call me at eight zero pm'
+
+    Args:
+        s: string to normalize
+        labels: labels used during model training.
+
+    Returns:
+            Normalized string
+    """
+
+    def good_token(token, labels):
+        s = set(labels)
+        for t in token:
+            if not t in s:
+                return False
+        return True
+
+    try:
+        text = _clean_text(s, ["english_cleaners"], table).strip()
+        return ''.join([t for t in text if good_token(t, labels=labels)])
+    except:
+        print("WARNING: Normalizing {} failed".format(s))
+        return None
+
+def ctc_decoder_predictions_tensor(prediction_cpu_tensor, batch_size, labels):
+    """
+    Takes output of greedy ctc decoder and performs ctc decoding algorithm to
+    remove duplicates and special symbol. Returns prediction
+    Args:
+        tensor: model output tensor
+        label: A list of labels
+    Returns:
+        prediction
+    """
+    blank_id = len(labels) - 1
+    hypotheses = []
+    labels_map = dict([(i, labels[i]) for i in range(len(labels))])
+    # iterate over batch
+    prediction_cpu_tensor = prediction_cpu_tensor.reshape((batch_size, int(prediction_cpu_tensor.size/batch_size)))
+    for ind in range(batch_size):
+        prediction = prediction_cpu_tensor[ind].tolist()
+        # CTC decoding procedure
+        decoded_prediction = []
+        previous = len(labels) - 1 # id of a blank symbol
+        for p in prediction:
+            if (p != previous or previous == blank_id) and p != blank_id:
+                decoded_prediction.append(p)
+            previous = p
+        hypothesis = ''.join([labels_map[c] for c in decoded_prediction])
+        hypotheses.append(hypothesis)
+    return hypotheses
+
+class SpeechClient(object):
+
+    def __init__(self, url, protocol, model_name, model_version, batch_size,
+                 model_platform=None, verbose=False,
+                 mode="batch",
+                 from_features=True):
+
+        self.model_name = model_name
+        self.model_version = model_version
+        self.verbose = verbose
+        self.batch_size = batch_size
+        self.transpose_audio_features = False
+        self.grpc_stub = None
+        self.ctx = None
+        self.correlation_id = 0
+        self.first_run = True
+        if mode == "streaming" or mode == "asynchronous":
+            self.correlation_id = 1
+
+        self.buffer = []
+
+        if protocol == "grpc":
+            # Create gRPC client for communicating with the server
+            self.prtcl_client = tritonclient.grpc
+        else:
+            # Create HTTP client for communicating with the server
+            self.prtcl_client = tritonclient.http
+
+        self.triton_client = self.prtcl_client.InferenceServerClient(
+            url=url, verbose=self.verbose)
+
+        self.audio_signals_name, self.num_samples_name, self.transcripts_name, \
+        self.audio_signals_type, self.num_samples_type, self.transcripts_type =  self.parse_model(# server_status,
+            model_name,
+            batch_size, model_platform, verbose)
+        self.labels = [" ", "a", "b", "c", "d", "e", "f", "g", "h", "i", "j", "k", "l", "m", "n", "o", "p", "q", "r", "s", "t", "u", "v", "w", "x", "y", "z", "'", "<BLANK>"]
+
+    def postprocess(self, transcript_values, labels):
+
+        res = []
+        for transcript, filename in zip(transcript_values,
+                                        labels):
+            print('---')
+            print('File: ', filename)
+            t=ctc_decoder_predictions_tensor(transcript, self.batch_size, self.labels)
+            print("Final transcript: ", t)
+            print('---')
+            res.append(t)
+        return res
+
+
+    def check_num_samples(self, num_samples, model_name):
+        if num_samples['data_type'] != 'TYPE_UINT32' and num_samples['data_type'] != 'TYPE_INT32':
+             raise Exception(
+                    "expecting num_samples datatype to be TYPE_UINT32/TYPE_INT32, "
+                    "model '" + model_name + "' output type is " +
+                    model_config.DataType.Name(num_samples['data_type']))
+        if len(num_samples['dims']) != 1:
+            raise Exception("Expecting num_samples to have 1 dimension, "
+                            "model '{}' num_samples has {}".format(
+                                model_name,len(num_samples['dims'])))
+
+    def parse_model(self, # server_status,
+                    model_name, batch_size,
+                    model_platform=None, verbose=False):
+        """
+        Check the configuration of the ensemble model
+        """
+
+        if self.prtcl_client is tritonclient.grpc:
+            config = self.triton_client.get_model_config(model_name, as_json=True)
+        else:
+            config = self.triton_client.get_model_config(model_name)
+
+        self.model_platform = model_platform
+
+        # Inputs are:
+        #   1) audio_signal: raw audio samples [num_samples]
+        #   2) sample_rate: sample rate of audio
+        #   3) num_samples: length of audio
+
+        if len(config['input']) < 2:
+            raise Exception(
+                "expecting 2-3 inputs, got {}".format(len(config['input'])))
+
+        # Outputs are:
+        #   1) transcripts:        candidate transcripts
+
+        if len(config['output']) != 1:
+            raise Exception(
+                "expecting 1 output, got {}".format(len(config['output'])))
+
+        audio_signal = config['input'][0]
+
+        if len(config['input']) > 1:
+            num_samples = config['input'][1]
+            self.check_num_samples(num_samples, model_name);
+
+        transcripts = config['output'][0]
+
+        expected_audio_signal_dim = 1
+
+        # Model specifying maximum batch size of 0 indicates that batching
+        # is not supported and so the input tensors do not expect an "N"
+        # dimension (and 'batch_size' should be 1 so that only a single
+        # image instance is inferred at a time).
+        max_batch_size = config['max_batch_size']
+        if max_batch_size == 0:
+            if batch_size != 1:
+                raise Exception(
+                    "batching not supported for model '" + model_name + "'")
+        else:  # max_batch_size > 0
+            if batch_size > max_batch_size:
+                raise Exception(
+                    "expecting batch size <= {} for model {}".format(
+                        max_batch_size, model_name))
+
+        if len(audio_signal['dims']) != expected_audio_signal_dim:
+            raise Exception("Expecting audio signal to have {} dimensions, "
+                            "model '{}' audio_signal has {}".format(
+                expected_audio_signal_dim,
+                model_name,
+                len(audio_signal.dims)))
+
+        return (audio_signal['name'],
+                num_samples['name'],
+                transcripts['name'],
+                triton_type_to_np_dtype[audio_signal['data_type']],
+                triton_type_to_np_dtype[num_samples['data_type']],
+                triton_type_to_np_dtype[transcripts['data_type']])
+
+
+    def recognize(self, audio_signal, filenames):
+        # Send requests of FLAGS.batch_size audio signals. If the number of
+        # audios isn't an exact multiple of FLAGS.batch_size then just
+        # start over with the first audio until the batch is filled.
+
+        input_batch = []
+        input_filenames = []
+        max_num_samples_batch = 0
+
+        for idx in range(self.batch_size):
+            input_batch.append(audio_signal[idx].astype(
+                self.audio_signals_type))
+            input_filenames.append(filenames[idx])
+            num_samples = audio_signal[idx].shape[0]
+
+            if (num_samples > max_num_samples_batch):
+                max_num_samples_batch = num_samples
+
+        for idx in range(self.batch_size):
+            num_samples = input_batch[idx].shape[0]
+
+            mean = np.mean(input_batch[idx])
+            std_var = np.std(input_batch[idx])
+            gauss_noise = np.random.normal(
+                mean,std_var,
+                max_num_samples_batch-num_samples)
+
+            input_batch[idx]= np.concatenate(
+                (input_batch[idx], gauss_noise.astype(
+                    self.audio_signals_type)))
+
+        max_num_samples_batch = np.asarray([max_num_samples_batch],
+                                           dtype=self.num_samples_type)
+
+        num_samples_batch = [max_num_samples_batch]*self.batch_size
+
+        # Send request
+        print("Sending request to transcribe file(s):", ",".join(
+            input_filenames))
+
+        inputs = []
+
+        input_batch = np.asarray(input_batch)
+        num_samples_batch = np.asarray(num_samples_batch)
+
+        inputs.append(self.prtcl_client.InferInput(self.audio_signals_name,
+                                                   input_batch.shape,
+                                                   np_to_triton_dtype(input_batch.dtype)))
+        inputs.append(self.prtcl_client.InferInput(self.num_samples_name,
+                                                   num_samples_batch.shape,
+                                                   np_to_triton_dtype(num_samples_batch.dtype)))
+
+        if self.prtcl_client is tritonclient.grpc:
+            inputs[0].set_data_from_numpy(input_batch)
+            inputs[1].set_data_from_numpy(num_samples_batch)
+        else: # http
+            inputs[0].set_data_from_numpy(input_batch, binary_data=True)
+            inputs[1].set_data_from_numpy(num_samples_batch, binary_data=True)
+
+        outputs = []
+        if self.prtcl_client is tritonclient.grpc:
+            outputs.append(self.prtcl_client.InferRequestedOutput(self.transcripts_name))
+        else:
+            outputs.append(self.prtcl_client.InferRequestedOutput(self.transcripts_name,
+                                                                  binary_data=True))
+
+        triton_result = self.triton_client.infer(self.model_name, inputs=inputs,
+                                          outputs=outputs)
+        transcripts = triton_result.as_numpy(self.transcripts_name)
+
+        result = self.postprocess(transcripts, input_filenames)
+
+        return result
+
+
+def preemphasis(signal, coeff=0.97):
+    return np.append(signal[0], signal[1:] - coeff*signal[:-1])
+
+
+def normalize_signal(signal, gain=None):
+    """
+    Normalize float32 signal to [-1, 1] range
+    """
+    if gain is None:
+        gain = 1.0/(np.max(np.abs(signal)) + 1e-5)
+    return signal*gain
+
+
+
+class AudioSegment(object):
+    """Monaural audio segment abstraction.
+    :param samples: Audio samples [num_samples x num_channels].
+    :type samples: ndarray.float32
+    :param sample_rate: Audio sample rate.
+    :type sample_rate: int
+    :raises TypeError: If the sample data type is not float or int.
+    """
+
+    def __init__(self, samples, sample_rate, target_sr=16000, trim=False,
+                 trim_db=60):
+        """Create audio segment from samples.
+        Samples are convert float32 internally, with int scaled to [-1, 1].
+        """
+        samples = self._convert_samples_to_float32(samples)
+        self._samples = samples
+        self._sample_rate = sample_rate
+        if self._samples.ndim >= 2:
+            self._samples = np.mean(self._samples, 1)
+
+    @staticmethod
+    def _convert_samples_to_float32(samples):
+        """Convert sample type to float32.
+        Audio sample type is usually integer or float-point.
+        Integers will be scaled to [-1, 1] in float32.
+        """
+        float32_samples = samples.astype('float32')
+        if samples.dtype in np.sctypes['int']:
+            bits = np.iinfo(samples.dtype).bits
+            float32_samples *= (1. / 2 ** (bits - 1))
+        elif samples.dtype in np.sctypes['float']:
+            pass
+        else:
+            raise TypeError("Unsupported sample type: %s." % samples.dtype)
+        return float32_samples
+
+    @classmethod
+    def from_file(cls, filename, target_sr=16000, int_values=False, offset=0,
+                  duration=0, trim=False):
+        """
+        Load a file supported by librosa and return as an AudioSegment.
+        :param filename: path of file to load
+        :param target_sr: the desired sample rate
+        :param int_values: if true, load samples as 32-bit integers
+        :param offset: offset in seconds when loading audio
+        :param duration: duration in seconds when loading audio
+        :return: numpy array of samples
+        """
+        with sf.SoundFile(filename, 'r') as f:
+            dtype = 'int32' if int_values else 'float32'
+            sample_rate = f.samplerate
+            if offset > 0:
+                f.seek(int(offset * sample_rate))
+            if duration > 0:
+                samples = f.read(int(duration * sample_rate), dtype=dtype)
+            else:
+                samples = f.read(dtype=dtype)
+
+        samples = samples.transpose()
+        return cls(samples, sample_rate, target_sr=target_sr, trim=trim)
+
+    @property
+    def samples(self):
+        return self._samples.copy()
+
+    @property
+    def sample_rate(self):
+        return self._sample_rate
+
+# define our clear function
+def clear_screen():
+    _ = system('clear')
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/tensorrt_io_props_fp16.json b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/tensorrt_io_props_fp16.json
new file mode 100644
index 0000000000..291f0f7dae
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/tensorrt_io_props_fp16.json
@@ -0,0 +1 @@
+{"input_shapes": [[-1, 64, -1]], "output_shapes": [[-1, -1, 29]], "input_types": ["TYPE_FP16"], "output_types": ["TYPE_FP16"], "input_names": ["input__0"], "output_names": ["output__0"], "dynamic_axes": {"input__0": [0, 2], "output__0": [0, 1]}, "min_shapes": [[1, 64, 200]], "opt_shapes": [[8, 64, 704]], "max_shapes": [[8, 64, 1680]]}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/tensorrt_io_props_fp32.json b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/tensorrt_io_props_fp32.json
new file mode 100644
index 0000000000..8a0fe1ab3d
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/tensorrt_io_props_fp32.json
@@ -0,0 +1 @@
+{"input_shapes": [[-1, 64, -1]], "output_shapes": [[-1, -1, 29]], "input_types": ["TYPE_FP32"], "output_types": ["TYPE_FP32"], "input_names": ["input__0"], "output_names": ["output__0"], "dynamic_axes": {"input__0": [0, 2], "output__0": [0, 1]}, "min_shapes": [[1, 64, 200]], "opt_shapes": [[8, 64, 704]], "max_shapes": [[8, 64, 1680]]}
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/triton_librispeech.csv b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/triton_librispeech.csv
new file mode 100644
index 0000000000..2ab5ed8477
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/triton/triton_librispeech.csv
@@ -0,0 +1,2 @@
+url,md5
+http://www.openslr.org/resources/12/test-clean.tar.gz,32fa31d27d2e1cad72775fee3f4849a9
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/__init__.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/__init__.py
new file mode 100644
index 0000000000..e69de29bb2
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/convert_librispeech.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/convert_librispeech.py
new file mode 100644
index 0000000000..40010110b5
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/convert_librispeech.py
@@ -0,0 +1,82 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+
+#!/usr/bin/env python
+import argparse
+import os
+import glob
+import multiprocessing
+import json
+
+import pandas as pd
+
+from preprocessing_utils import parallel_preprocess
+
+parser = argparse.ArgumentParser(description='Preprocess LibriSpeech.')
+parser.add_argument('--input_dir', type=str, required=True,
+                    help='LibriSpeech collection input dir')
+parser.add_argument('--dest_dir', type=str, required=True,
+                    help='Output dir')
+parser.add_argument('--output_json', type=str, default='./',
+                    help='name of the output json file.')
+parser.add_argument('-s','--speed', type=float, nargs='*',
+                    help='Speed perturbation ratio')
+parser.add_argument('--target_sr', type=int, default=None,
+                    help='Target sample rate. '
+                         'defaults to the input sample rate')
+parser.add_argument('--overwrite', action='store_true',
+                    help='Overwrite file if exists')
+parser.add_argument('--parallel', type=int, default=multiprocessing.cpu_count(),
+                    help='Number of threads to use when processing audio files')
+args = parser.parse_args()
+
+args.input_dir = args.input_dir.rstrip('/')
+args.dest_dir = args.dest_dir.rstrip('/')
+
+def build_input_arr(input_dir):
+    txt_files = glob.glob(os.path.join(input_dir, '**', '*.trans.txt'),
+                          recursive=True)
+
+    input_data = []
+    for txt_file in txt_files:
+        rel_path = os.path.relpath(txt_file, input_dir)
+        with open(txt_file) as fp:
+            for line in fp:
+                fname, _, transcript = line.partition(' ')
+                input_data.append(dict(input_relpath=os.path.dirname(rel_path),
+                                       input_fname=fname+'.flac',
+                                       transcript=transcript))
+    return input_data
+
+
+print("[%s] Scaning input dir..." % args.output_json)
+dataset = build_input_arr(input_dir=args.input_dir)
+
+print("[%s] Converting audio files..." % args.output_json)
+dataset = parallel_preprocess(dataset=dataset,
+                              input_dir=args.input_dir,
+                              dest_dir=args.dest_dir,
+                              target_sr=args.target_sr,
+                              speed=args.speed,
+                              overwrite=args.overwrite,
+                              parallel=args.parallel)
+
+print("[%s] Generating json..." % args.output_json)
+df = pd.DataFrame(dataset, dtype=object)
+
+# Save json with python. df.to_json() produces back slashed in file paths
+dataset = df.to_dict(orient='records')
+with open(args.output_json, 'w') as fp:
+    json.dump(dataset, fp, indent=2)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/download_librispeech.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/download_librispeech.py
new file mode 100644
index 0000000000..ad36ad4e42
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/download_librispeech.py
@@ -0,0 +1,72 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+#!/usr/bin/env python
+
+import os
+import argparse
+import pandas as pd
+
+from download_utils import download_file, md5_checksum, extract
+
+parser = argparse.ArgumentParser(description='Download, verify and extract dataset files')
+parser.add_argument('csv', type=str,
+                    help='CSV file with urls and checksums to download.')
+parser.add_argument('dest', type=str,
+                    help='Download destnation folder.')
+parser.add_argument('-e', type=str, default=None,
+                    help='Extraction destnation folder. Defaults to download folder if not provided')
+parser.add_argument('--skip_download', action='store_true',
+                    help='Skip downloading the files')
+parser.add_argument('--skip_checksum', action='store_true',
+                    help='Skip checksum')
+parser.add_argument('--skip_extract', action='store_true',
+                    help='Skip extracting files')
+args = parser.parse_args()
+args.e = args.e or args.dest
+
+
+df = pd.read_csv(args.csv, delimiter=',')
+
+
+if not args.skip_download:
+    for url in df.url:
+        fname = url.split('/')[-1]
+        print("Downloading %s:" % fname)
+        download_file(url=url, dest_folder=args.dest, fname=fname)
+else:
+    print("Skipping file download")
+
+
+if not args.skip_checksum:
+    for index, row in df.iterrows():
+        url = row['url']
+        md5 = row['md5']
+        fname = url.split('/')[-1]
+        fpath = os.path.join(args.dest, fname)
+        print("Verifing %s: " % fname, end='')
+        ret = md5_checksum(fpath=fpath, target_hash=md5)
+        print("Passed" if ret else "Failed")
+else:
+    print("Skipping checksum")
+
+
+if not args.skip_extract:
+    for url in df.url:
+        fname = url.split('/')[-1]
+        fpath = os.path.join(args.dest, fname)
+        print("Decompressing %s:" % fpath)
+        extract(fpath=fpath, dest_folder=args.e)
+else:
+    print("Skipping file extraction")
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/download_utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/download_utils.py
new file mode 100644
index 0000000000..fd6644404b
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/download_utils.py
@@ -0,0 +1,71 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+#!/usr/bin/env python
+
+import hashlib
+import requests
+import os
+import tarfile
+import tqdm
+
+def download_file(url, dest_folder, fname, overwrite=False):
+    fpath = os.path.join(dest_folder, fname)
+    if os.path.isfile(fpath):
+        if overwrite:
+            print("Overwriting existing file")
+        else:
+            print("File exists, skipping download.")
+            return
+
+    tmp_fpath = fpath + '.tmp'
+
+    if not os.path.exists(os.path.dirname(tmp_fpath)):
+        os.makedirs(os.path.dirname(tmp_fpath))
+
+    r = requests.get(url, stream=True)
+    file_size = int(r.headers['Content-Length'])
+    chunk_size = 1024 * 1024  # 1MB
+    total_chunks = int(file_size / chunk_size)
+
+    with open(tmp_fpath, 'wb') as fp:
+        content_iterator = r.iter_content(chunk_size=chunk_size)
+        chunks = tqdm.tqdm(content_iterator, total=total_chunks,
+                           unit='MB', desc=fpath, leave=True)
+        for chunk in chunks:
+            fp.write(chunk)
+
+    os.rename(tmp_fpath, fpath)
+
+
+def md5_checksum(fpath, target_hash):
+    file_hash = hashlib.md5()
+    with open(fpath, "rb") as fp:
+        for chunk in iter(lambda: fp.read(1024*1024), b""):
+            file_hash.update(chunk)
+    return file_hash.hexdigest() == target_hash
+
+
+def extract(fpath, dest_folder):
+    if fpath.endswith('.tar.gz'):
+        mode = 'r:gz'
+    elif fpath.endswith('.tar'):
+        mode = 'r:'
+    else:
+        raise IOError('fpath has unknown extention: %s' % fpath)
+
+    with tarfile.open(fpath, mode) as tar:
+        members = tar.getmembers()
+        for member in tqdm.tqdm(iterable=members, total=len(members), leave=True):
+            tar.extract(path=dest_folder, member=member)
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/inference_librispeech.csv b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/inference_librispeech.csv
new file mode 100644
index 0000000000..40dac4e0e6
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/inference_librispeech.csv
@@ -0,0 +1,5 @@
+url,md5
+http://www.openslr.org/resources/12/dev-clean.tar.gz,42e2234ba48799c1f50f24a7926300a1
+http://www.openslr.org/resources/12/dev-other.tar.gz,c8d0bcc9cca99d4f8b62fcc847357931
+http://www.openslr.org/resources/12/test-clean.tar.gz,32fa31d27d2e1cad72775fee3f4849a9
+http://www.openslr.org/resources/12/test-other.tar.gz,fb5a50374b501bb3bac4815ee91d3135
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/librispeech.csv b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/librispeech.csv
new file mode 100644
index 0000000000..d48a9f8db7
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/librispeech.csv
@@ -0,0 +1,8 @@
+url,md5
+http://www.openslr.org/resources/12/dev-clean.tar.gz,42e2234ba48799c1f50f24a7926300a1
+http://www.openslr.org/resources/12/dev-other.tar.gz,c8d0bcc9cca99d4f8b62fcc847357931
+http://www.openslr.org/resources/12/test-clean.tar.gz,32fa31d27d2e1cad72775fee3f4849a9
+http://www.openslr.org/resources/12/test-other.tar.gz,fb5a50374b501bb3bac4815ee91d3135
+http://www.openslr.org/resources/12/train-clean-100.tar.gz,2a93770f6d5c6c964bc36631d331a522
+http://www.openslr.org/resources/12/train-clean-360.tar.gz,c0e676e450a7ff2f54aeade5171606fa
+http://www.openslr.org/resources/12/train-other-500.tar.gz,d1a0fd59409feb2c614ce4d30c387708
diff --git a/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/preprocessing_utils.py b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/preprocessing_utils.py
new file mode 100644
index 0000000000..15605cea27
--- /dev/null
+++ b/PyTorch/contrib/audio/jasper/Jasper_pytorch_wuqingdian01/utils/preprocessing_utils.py
@@ -0,0 +1,76 @@
+# Copyright (c) 2019, NVIDIA CORPORATION. All rights reserved.
+#
+# 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.
+
+#!/usr/bin/env python
+import os
+import multiprocessing
+import librosa
+import functools
+
+import sox
+
+
+from tqdm import tqdm
+
+def preprocess(data, input_dir, dest_dir, target_sr=None, speed=None,
+               overwrite=True):
+    speed = speed or []
+    speed.append(1)
+    speed = list(set(speed))  # Make uniqe
+
+    input_fname = os.path.join(input_dir,
+                               data['input_relpath'],
+                               data['input_fname'])
+    input_sr = sox.file_info.sample_rate(input_fname)
+    target_sr = target_sr or input_sr
+
+    os.makedirs(os.path.join(dest_dir, data['input_relpath']), exist_ok=True)
+
+    output_dict = {}
+    output_dict['transcript'] = data['transcript'].lower().strip()
+    output_dict['files'] = []
+
+    fname = os.path.splitext(data['input_fname'])[0]
+    for s in speed:
+        output_fname = fname + '{}.wav'.format('' if s==1 else '-{}'.format(s))
+        output_fpath = os.path.join(dest_dir,
+                                    data['input_relpath'],
+                                    output_fname)
+
+        if not os.path.exists(output_fpath) or overwrite:
+            cbn = sox.Transformer().speed(factor=s).convert(target_sr)
+            cbn.build(input_fname, output_fpath)
+
+        file_info = sox.file_info.info(output_fpath)
+        file_info['fname'] = os.path.join(os.path.basename(dest_dir),
+                                          data['input_relpath'],
+                                          output_fname)
+        file_info['speed'] = s
+        output_dict['files'].append(file_info)
+
+        if s == 1:
+            file_info = sox.file_info.info(output_fpath)
+            output_dict['original_duration'] = file_info['duration']
+            output_dict['original_num_samples'] = file_info['num_samples']
+
+    return output_dict
+
+
+def parallel_preprocess(dataset, input_dir, dest_dir, target_sr, speed, overwrite, parallel):
+    with multiprocessing.Pool(parallel) as p:
+        func = functools.partial(preprocess,
+            input_dir=input_dir, dest_dir=dest_dir,
+            target_sr=target_sr, speed=speed, overwrite=overwrite)
+        dataset = list(tqdm(p.imap(func, dataset), total=len(dataset)))
+        return dataset
-- 
Gitee