Searched refs:tf_axis (Results 1 – 2 of 2) sorted by relevance
| /external/tensorflow/tensorflow/python/kernel_tests/array_ops/ |
| D | gather_op_test.py | 127 tf_axis = constant_op.constant(axis)
129 gather = array_ops.gather(tf_params, tf_indices, axis=tf_axis)
148 gather, [tf_params, tf_indices, tf_axis], gather_grad)
193 tf_axis = constant_op.constant(axis)
196 tape.watch(tf_axis)
198 gather = array_ops.gather(tf_params, tf_indices, axis=tf_axis)
217 gather, [tf_params, tf_indices, tf_axis], gather_grad)
|
| /external/tensorflow/tensorflow/compiler/tf2tensorrt/convert/ |
| D | convert_nodes.cc | 529 Status ConvertAxis(int tf_axis, int trt_nb_dims, absl::string_view node_name, in ConvertAxis() argument
533 if (tf_axis < -tf_nb_dims || tf_axis >= tf_nb_dims) { in ConvertAxis()
535 "Axis value of ", tf_axis, " is out of bounds, must be in range [", in ConvertAxis()
539 if (tf_axis < 0) tf_axis += tf_nb_dims; in ConvertAxis()
541 if (use_implicit_batch && tf_axis == 0) { in ConvertAxis()
547 *trt_axis = use_implicit_batch ? tf_axis - 1 : tf_axis; in ConvertAxis()
2816 for (int tf_axis : squeeze_dims) { in ConvertSqueeze() local
2820 TF_RETURN_IF_ERROR(ConvertAxis(tf_axis, dims.nbDims, node_def.name(), in ConvertSqueeze()
2825 "Dimension ", tf_axis, " with size ", input_dims[trt_axis], in ConvertSqueeze()
4768 const int64 tf_axis = attrs.get<int64>("axis"); in ConvertPack() local
[all …]
|