| 1400 | } |
| 1401 | |
| 1402 | void Compute(OpKernelContext* ctx) override NO_THREAD_SAFETY_ANALYSIS { |
| 1403 | auto locks = MaybeLockEmbeddingVariableInputMutexesInOrder<Tindex, T>(ctx, use_exclusive_lock_, |
| 1404 | {0, 1, 2}); |
| 1405 | EmbeddingVar<Tindex, T>* var = nullptr; |
| 1406 | OP_REQUIRES_OK(ctx, GetInputEmbeddingVar(ctx, 0, &var)); |
| 1407 | core::ScopedUnref unref_var(var); |
| 1408 | |
| 1409 | EmbeddingVar<Tindex, T>* m = nullptr; |
| 1410 | OP_REQUIRES_OK(ctx, GetInputEmbeddingVar(ctx, 1, &m)); |
| 1411 | core::ScopedUnref unref_m(m); |
| 1412 | |
| 1413 | EmbeddingVar<Tindex, T>* v = nullptr; |
| 1414 | OP_REQUIRES_OK(ctx, GetInputEmbeddingVar(ctx, 2, &v)); |
| 1415 | core::ScopedUnref unref_v(v); |
| 1416 | |
| 1417 | const Tensor& beta1_power = ctx->input(3); |
| 1418 | const Tensor& beta2_power = ctx->input(4); |
| 1419 | const Tensor& lr = ctx->input(5); |
| 1420 | const Tensor& beta1 = ctx->input(6); |
| 1421 | const Tensor& beta2 = ctx->input(7); |
| 1422 | const Tensor& epsilon = ctx->input(8); |
| 1423 | const Tensor& grad = ctx->input(9); |
| 1424 | const Tensor& indices = ctx->input(10); |
| 1425 | const Tensor& global_step = ctx->input(11); |
| 1426 | |
| 1427 | OP_REQUIRES( |
| 1428 | ctx, TensorShapeUtils::IsScalar(beta1_power.shape()), |
| 1429 | errors::InvalidArgument("beta1_power is not a scalar: ", |
| 1430 | beta1_power.shape().DebugString())); |
| 1431 | OP_REQUIRES( |
| 1432 | ctx, TensorShapeUtils::IsScalar(beta2_power.shape()), |
| 1433 | errors::InvalidArgument("beta2_power is not a scalar: ", |
| 1434 | beta2_power.shape().DebugString())); |
| 1435 | OP_REQUIRES( |
| 1436 | ctx, TensorShapeUtils::IsScalar(lr.shape()), |
| 1437 | errors::InvalidArgument("lr is not a scalar: ", |
| 1438 | lr.shape().DebugString())); |
| 1439 | OP_REQUIRES( |
| 1440 | ctx, TensorShapeUtils::IsScalar(beta1.shape()), |
| 1441 | errors::InvalidArgument("beta1 is not a scalar: ", |
| 1442 | beta1.shape().DebugString())); |
| 1443 | OP_REQUIRES( |
| 1444 | ctx, TensorShapeUtils::IsScalar(beta2.shape()), |
| 1445 | errors::InvalidArgument("beta2 is not a scalar: ", |
| 1446 | beta2.shape().DebugString())); |
| 1447 | OP_REQUIRES( |
| 1448 | ctx, TensorShapeUtils::IsScalar(epsilon.shape()), |
| 1449 | errors::InvalidArgument("epsilon is not a scalar: ", |
| 1450 | epsilon.shape().DebugString())); |
| 1451 | OP_REQUIRES( |
| 1452 | ctx, TensorShapeUtils::IsVector(indices.shape()), |
| 1453 | errors::InvalidArgument("indices must be one-dimensional")); |
| 1454 | |
| 1455 | int64 inner_dim = 1; |
| 1456 | TensorShape var_shape({var->ValueLen()}); |
| 1457 | for (int d = 0; d < var_shape.dims(); d++) { |
| 1458 | OP_REQUIRES(ctx, var_shape.dim_size(d) == grad.dim_size(d + 1), |
| 1459 | errors::InvalidArgument(strings::StrCat( |
nothing calls this directly
no test coverage detected