void Compute()

in src/optimize_op.cc [66:112]
37 lines of code
3 McCabe index (conditional complexity)

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

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

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

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

    Tensor param = ctx->mutable_input(0, false);
    Tensor cv    = ctx->mutable_input(1, true);
    Tensor rv    = ctx->mutable_input(2, false);

    OP_REQUIRES(ctx, param.dims() == 2, errors::InvalidArgument("only applies to 2d params"));

    uint C = param.dim_size(0);
    uint K = param.dim_size(1);

    OP_REQUIRES(ctx, cv.shape().num_elements() == K, errors::InvalidArgument("bad cv shape"));
    OP_REQUIRES(ctx, rv.shape().num_elements() == C, errors::InvalidArgument("bad rv shape"));

    Tensor* x; OP_REQUIRES_OK(ctx, ctx->allocate_output(3,  param.shape(),      &x));
    Tensor* m; OP_REQUIRES_OK(ctx, ctx->allocate_output(4,  TensorShape({ 2 }), &m));

    CUstream stream = ((CUDAStream*)ctx->op_device_context()->stream()->implementation())->cuda_stream();

    Adafactor<V1,V4>(stream, SMs_,
      cv.flat<float>().data(),
      rv.flat<float>().data(),
      x->flat<float>().data(),
      m->flat<float>().data(),
      param.flat<float>().data(),
      (const V1*)grad.flat<T>().data(),
      norm_scale_ptr,
      scal.scalar<float>()(), lr.scalar<float>()(), decay.scalar<float>()(), epsilon_, clip.scalar<float>()(), C, K, saturate_, zero_infs_, zero_nans_
    );
  }