void Compute()

in src/optimize_op.cc [368:429]
54 lines of code
5 McCabe index (conditional complexity)

  void Compute(OpKernelContext* ctx) override
  {
    if (SMs_ == 0)
      SMs_ = GetCountSMs();

    ctx->forward_ref_input_to_ref_output(1, 0);
    ctx->forward_ref_input_to_ref_output(2, 1);
    ctx->forward_ref_input_to_ref_output(3, 2);

    const Tensor& grad = ctx->input(0);
    const Tensor& lr   = ctx->input(4);
    const Tensor& scal = ctx->input(5);
    const Tensor& clip = ctx->input(6);

    OpInputList norm_scale, gate;
    ctx->input_list("norm_scale", &norm_scale);
    ctx->input_list("gate", &gate);
    const float* norm_scale_ptr = norm_scale.size() > 0 ? norm_scale[0].flat<float>().data() : NULL;

    Tensor param = ctx->mutable_input(1, false);
    Tensor mean  = ctx->mutable_input(2, false);
    Tensor var   = ctx->mutable_input(3, false);

    uint size, K;
    if (lazy_emb_)
    {
      OP_REQUIRES(ctx, param.dims() == 2, errors::InvalidArgument("lazy_emb only applies to 2d embedding params"));
      size = param.dim_size(0);
      K    = param.dim_size(1);
    }
    else
    {
      size = param.shape().num_elements();
      K = 0;
    }
    CUstream stream = ((CUDAStream*)ctx->op_device_context()->stream()->implementation())->cuda_stream();

    if (gate.size() > 0)
    {
      uint blocks = param.dim_size(0);
      uint bsize  = param.dim_size(1);

      ApplyAdamGated<VG,VRM,VRV>(stream,
        gate[0].flat<float>().data(),
        (const VG*)grad.flat<TG>().data(),
        norm_scale_ptr,
        param.flat<float>().data(),
        (VRM*)mean.flat<TR>().data(),
        (VRV*)var.flat<TR>().data(),
        lr.scalar<float>()(), decay_mean_, decay_var_, epsilon_, scal.scalar<float>()(), clip.scalar<float>()(), blocks, bsize, saturate_, zero_infs_, zero_nans_
      );
    }

    ApplyAdam<VG,VRM,VRV>(stream, SMs_,
      (const VG*)grad.flat<TG>().data(),
      norm_scale_ptr,
      param.flat<float>().data(),
      (VRM*)mean.flat<TR>().data(),
      (VRV*)var.flat<TR>().data(),
      lr.scalar<float>()(), decay_mean_, decay_var_, epsilon_, scal.scalar<float>()(), clip.scalar<float>()(), size, K, saturate_, zero_infs_, zero_nans_
    );
  }