Searched refs:axis_input (Results 1 – 4 of 4) sorted by relevance
| /external/tensorflow/tensorflow/core/kernels/ |
| D | sparse_split_op.cc | 33 const int64 axis_input = context->input(0).scalar<int64>()(); in Compute() local
52 const int64 axis = (axis_input < 0) ? input_rank + axis_input : axis_input; in Compute()
57 ", ", input_rank, "), got ", axis_input)); in Compute()
|
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | gather_op.cc | 173 int64 axis_input; in XlaGatherWithBatchDimsOpImpl() local
174 TF_RETURN_IF_ERROR(context->ConstantInputAsIntScalar(2, &axis_input)); in XlaGatherWithBatchDimsOpImpl()
177 if (-params_dims > axis_input || axis_input >= params_dims) { in XlaGatherWithBatchDimsOpImpl()
180 "), but got ", axis_input); in XlaGatherWithBatchDimsOpImpl()
182 if (axis_input < 0) { in XlaGatherWithBatchDimsOpImpl()
183 axis_input += params_dims; in XlaGatherWithBatchDimsOpImpl()
185 axis = axis_input; in XlaGatherWithBatchDimsOpImpl()
|
| /external/tensorflow/tensorflow/core/grappler/optimizers/ |
| D | arithmetic_optimizer.cc | 3588 bool IsAxis0(const NodeDef& node, int axis_input) { in IsAxis0() argument
3590 if (!GetTensorFromConstNode(node.input(axis_input), &axis_tensor)) in IsAxis0()
|
| /external/tensorflow/tensorflow/compiler/tf2tensorrt/convert/ |
| D | convert_nodes.cc | 5320 const auto& axis_input = inputs.at(2); in ConvertGather() local
5328 absl::Span<const int> axis = axis_input.weights().GetSpan<int>(); in ConvertGather()
|