Home
last modified time | relevance | path

Searched refs:axis_tensor (Results 1 – 5 of 5) sorted by relevance

/external/tensorflow/tensorflow/core/kernels/
Dunique_op.cc56 const Tensor& axis_tensor = context->input(1); in Compute() local
57 OP_REQUIRES(context, TensorShapeUtils::IsVector(axis_tensor.shape()), in Compute()
60 context, axis_tensor.NumElements() <= 1, in Compute()
63 if (axis_tensor.NumElements() == 0) { in Compute()
68 (axis_tensor.dtype() == DT_INT32 || in Compute()
69 axis_tensor.dtype() == DT_INT64), in Compute()
72 DataTypeString(axis_tensor.dtype()))); in Compute()
73 if (axis_tensor.dtype() == DT_INT32) { in Compute()
74 axis = internal::SubtleMustCopy(axis_tensor.scalar<int32>()()); in Compute()
76 axis = internal::SubtleMustCopy(axis_tensor.scalar<int64>()()); in Compute()
Dgather_op.cc55 const Tensor& axis_tensor = c->input(2); in Compute() local
56 OP_REQUIRES(c, TensorShapeUtils::IsScalar(axis_tensor.shape()), in Compute()
59 if (axis_tensor.dtype() == DT_INT32) { in Compute()
60 axis = axis_tensor.scalar<int32>()(); in Compute()
61 } else if (axis_tensor.dtype() == DT_INT64) { in Compute()
62 axis = axis_tensor.scalar<int64>()(); in Compute()
/external/tensorflow/tensorflow/lite/kernels/
Dreverse.cc69 const TfLiteTensor* axis_tensor = GetInput(context, node, kAxisTensor); in Eval() local
70 int axis = GetTensorData<int32_t>(axis_tensor)[0]; in Eval()
/external/tensorflow/tensorflow/core/ops/
Darray_ops.cc936 const Tensor* axis_tensor = c->input_tensor(1); in __anon7c94107b1402() local
937 if (axis_tensor != nullptr && c->RankKnown(input)) { in __anon7c94107b1402()
940 if (axis_tensor->dtype() == DT_INT32) { in __anon7c94107b1402()
941 axis_value = AsInt64<int32>(axis_tensor, axis_tensor->NumElements()); in __anon7c94107b1402()
943 axis_value = AsInt64<int64>(axis_tensor, axis_tensor->NumElements()); in __anon7c94107b1402()
/external/tensorflow/tensorflow/core/grappler/optimizers/
Dconstant_folding.cc156 Tensor axis_tensor(axis_input.dtype(), axis_input.shape()); in GetConcatAxis() local
157 if (!axis_tensor.FromProto(axis_input.value())) { in GetConcatAxis()
161 ? static_cast<int>(axis_tensor.scalar<int64>()()) in GetConcatAxis()
162 : axis_tensor.scalar<int32>()(); in GetConcatAxis()