TF_RETURN_IF_ERROR(forward_attrs_.Get("data_format", &data_format));

    // Grad for A
    grad_inputs[0] = upstream_grad;
    grad_inputs[0]->Ref();

    // Grad for bias
    std::string name = "bias_add_grad";
    TF_RETURN_IF_ERROR(BiasAddGrad(ctx, upstream_grad, &grad_inputs[1],
                                   data_format.c_str(), name.c_str()));

    return absl::OkStatus();
  }

    A.reset(A_raw);
  }
  // Bias
  float Bias_vals[] = {2.0f, 3.0f};
  int64_t Bias_dims[] = {2};
  AbstractTensorHandlePtr Bias;
  {
    AbstractTensorHandle* Bias_raw;
    status_ = TestTensorHandleWithDims<float, TF_FLOAT>(
        immediate_execution_ctx_.get(), Bias_vals, Bias_dims, 1, &Bias_raw);
    ASSERT_EQ(errors::OK, status_.code()) << status_.message();
    Bias.reset(Bias_raw);
  }

    *   XLA reduction emitter is deterministic when the environment variable
        `TF_DETERMINISTIC_OPS` is set to "true" or "1". This extends
        deterministic `tf.nn.bias_add` back-prop functionality (and therefore
        also deterministic back-prop of bias-addition in Keras layers) to
        include when XLA JIT compilation is enabled.
    *   Fix problem, when running on a CUDA GPU and when either environment