| 435 | } |
| 436 | |
| 437 | void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS { |
| 438 | auto locks = |
| 439 | MaybeLockEmbeddingVariableInputMutexesInOrder<TKey, T>(ctx, use_exclusive_lock_, {0, 1, 2}); |
| 440 | |
| 441 | EmbeddingVar<TKey, T>* var_ = nullptr; |
| 442 | OP_REQUIRES_OK(ctx, GetInputEmbeddingVar(ctx, 0, &var_)); |
| 443 | core::ScopedUnref unref_var(var_); |
| 444 | EmbeddingVar<TKey, T>* accum_ = nullptr; |
| 445 | OP_REQUIRES_OK(ctx, GetInputEmbeddingVar(ctx, 1, &accum_)); |
| 446 | core::ScopedUnref unref_accum(accum_); |
| 447 | EmbeddingVar<TKey, T>* linear_ = nullptr; |
| 448 | OP_REQUIRES_OK(ctx, GetInputEmbeddingVar(ctx, 2, &linear_)); |
| 449 | core::ScopedUnref unref_linear(linear_); |
| 450 | |
| 451 | const Tensor& grad = ctx->input(3); |
| 452 | const Tensor& indices = ctx->input(4); |
| 453 | OP_REQUIRES(ctx, TensorShapeUtils::IsVector(indices.shape()), |
| 454 | errors::InvalidArgument("indices must be one-dimensional")); |
| 455 | |
| 456 | const Tensor& lr = ctx->input(5); |
| 457 | OP_REQUIRES(ctx, |
| 458 | TensorShapeUtils::IsScalar(lr.shape()) && |
| 459 | lr.scalar<T>()() > static_cast<T>(0), |
| 460 | errors::InvalidArgument("lr is not a positive scalar: ", |
| 461 | lr.shape().DebugString())); |
| 462 | |
| 463 | const Tensor& l1 = ctx->input(6); |
| 464 | OP_REQUIRES(ctx, |
| 465 | TensorShapeUtils::IsScalar(l1.shape()) && |
| 466 | l1.scalar<T>()() >= static_cast<T>(0), |
| 467 | errors::InvalidArgument("l1 regularization strength is not a " |
| 468 | "non-negative scalar: ", |
| 469 | l1.shape().DebugString())); |
| 470 | const Tensor& l2 = ctx->input(7); |
| 471 | OP_REQUIRES(ctx, |
| 472 | TensorShapeUtils::IsScalar(l2.shape()) && |
| 473 | l2.scalar<T>()() >= static_cast<T>(0), |
| 474 | errors::InvalidArgument("l2 regularization strength is not a " |
| 475 | "non-negative scalar: ", |
| 476 | l2.shape().DebugString())); |
| 477 | const int lr_power_index = has_l2_shrinkage ? 9 : 8; |
| 478 | const Tensor& lr_power = ctx->input(lr_power_index); |
| 479 | OP_REQUIRES(ctx, |
| 480 | TensorShapeUtils::IsScalar(lr_power.shape()) && |
| 481 | lr_power.scalar<T>()() <= static_cast<T>(0), |
| 482 | errors::InvalidArgument("lr_power is not a " |
| 483 | "non-positive scalar: ", |
| 484 | lr_power.shape().DebugString())); |
| 485 | int64 inner_dim = 1; |
| 486 | TensorShape var_shape({var_->ValueLen()}); |
| 487 | for (int d = 0; d < var_shape.dims(); d++) { |
| 488 | OP_REQUIRES(ctx, var_shape.dim_size(d) == grad.dim_size(d + 1), |
| 489 | errors::InvalidArgument(strings::StrCat( |
| 490 | "var and grad must match in dimension ", d + 1))); |
| 491 | inner_dim *= grad.dim_size(d + 1); |
| 492 | } |
| 493 | const int64 N = indices.dim_size(0); |
| 494 | OP_REQUIRES( |
nothing calls this directly
no test coverage detected