Searched refs:last_dim_size (Results 1 – 6 of 6) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | topk_op.cc | 43 int last_dim_size = input_shape.dim_size(last_dim); in Compile() local
45 context, last_dim_size >= k, in Compile()
47 last_dim_size, ", needed ", k)); in Compile()
48 if (last_dim_size < k) { in Compile()
49 k = last_dim_size; in Compile()
|
| D | diag_op.cc | 33 xla::XlaOp CreateDiagonal(xla::XlaOp input, int64 last_dim_size, in CreateDiagonal() argument
45 xla::XlaOp iota = xla::Iota(builder, xla::S32, last_dim_size); in CreateDiagonal()
46 xla::XlaOp iota_broadcast = xla::Broadcast(iota, {last_dim_size}); in CreateDiagonal()
63 out_dim_sizes.push_back(last_dim_size); in CreateDiagonal()
64 out_dim_sizes.push_back(last_dim_size); in CreateDiagonal()
|
| /external/tensorflow/tensorflow/compiler/xla/client/lib/ |
| D | sorting.cc | 33 int64 last_dim_size = input_shape.dimensions(last_dim); in TopK() local
40 (kPerPartitionSize / k > 2) && last_dim_size >= kLastDimSizeThreshold) { in TopK()
42 CeilOfRatio(last_dim_size - k, kPerPartitionSize - k); in TopK()
83 int64 last_dim_size = input_shape.dimensions(last_dim); in TopKWithPartitions() local
84 const int64 per_partition_size = CeilOfRatio(last_dim_size, num_partitions); in TopKWithPartitions()
|
| /external/tensorflow/tensorflow/core/kernels/ |
| D | decode_proto_op.cc | 890 last_dim_size = tensor->dim_size(tensor->dims() - 1); in AccumulateFields()
895 stride = last_dim_size * element_size; in AccumulateFields()
901 stride = last_dim_size * sizeof(tstring); in AccumulateFields()
907 int last_dim_size; in AccumulateFields() member
931 default_value, info.last_dim_size); in AccumulateFields()
|
| /external/tensorflow/tensorflow/python/ops/ |
| D | nn_ops.py | 3553 last_dim_size = array_ops.slice(
3555 output = array_ops.reshape(logits, array_ops.concat([[-1], last_dim_size], 0))
|
| /external/tensorflow/tensorflow/compiler/mlir/xla/transforms/ |
| D | legalize_tf.cc | 4746 int64_t last_dim_size = input_type.getDimSize(last_dim_index); in matchAndRewrite() local
4747 if (last_dim_size == ShapedType::kDynamicSize) return failure(); in matchAndRewrite()
4777 std::min((*k_attr.begin()).getSExtValue(), last_dim_size); in matchAndRewrite()
|