Searched refs:dimno (Results 1 – 7 of 7) sorted by relevance
| /external/tensorflow/tensorflow/compiler/tf2xla/kernels/ |
| D | mirror_pad_op.cc | 43 for (int64_t dimno = original_shape.rank() - 1; dimno >= 0; --dimno) { in DoMirrorPad() local
44 auto t_rev = xla::Rev(accum, {dimno}); in DoMirrorPad()
45 int64_t lhs_padding = pad_literal.Get<int64_t>({dimno, 0}); in DoMirrorPad()
46 int64_t rhs_padding = pad_literal.Get<int64_t>({dimno, 1}); in DoMirrorPad()
47 int64_t dim_size = original_shape.dimensions(dimno); in DoMirrorPad()
58 dim_size - excluded_edges, 1, dimno); in DoMirrorPad()
60 excluded_edges + rhs_padding, 1, dimno); in DoMirrorPad()
61 accum = xla::ConcatInDim(b, {lhs_pad, accum, rhs_pad}, dimno); in DoMirrorPad()
130 for (int64_t dimno = original_shape.rank() - 1; dimno >= 0; --dimno) { in DoMirrorPadGrad() local
131 int64_t lhs_padding = pad_literal.Get<int64_t>({dimno, 0}); in DoMirrorPadGrad()
[all …]
|
| /external/tensorflow/tensorflow/compiler/xla/ |
| D | index_util.cc | 118 for (int64_t dimno = indices.size() - 1; dimno >= 0; --dimno) { in BumpIndices() local
119 int64_t limit = shape.dimensions(dimno); in BumpIndices()
120 if (indices[dimno] + 1 < limit) { in BumpIndices()
121 indices[dimno]++; in BumpIndices()
124 std::fill(indices.begin() + dimno + 1, indices.end(), 0); in BumpIndices()
|
| /external/tensorflow/tensorflow/compiler/xla/pjrt/ |
| D | transpose_test.cc | 148 for (int dimno = indices.size() - 1; dimno >= 0; --dimno) { in BumpIndices() local
149 if (indices[dimno] + 1 < shape[dimno]) { in BumpIndices()
150 indices[dimno]++; in BumpIndices()
153 std::fill(indices.begin() + dimno + 1, indices.end(), 0); in BumpIndices()
|
| /external/tensorflow/tensorflow/compiler/xla/client/ |
| D | xla_builder.h | 464 XlaOp PadInDim(XlaOp operand, XlaOp padding_value, int64_t dimno,
495 int64_t limit_index, int64_t stride, int64_t dimno);
1114 friend XlaOp PadInDim(XlaOp operand, XlaOp padding_value, int64_t dimno,
1139 int64_t limit_index, int64_t stride, int64_t dimno);
1779 XlaOp PadInDim(XlaOp operand, XlaOp padding_value, int64_t dimno,
1855 int64_t stride, int64_t dimno);
|
| D | xla_builder.cc | 1047 int64_t dimno) { in SliceInDim() argument
1054 starts[dimno] = start_index; in SliceInDim()
1055 limits[dimno] = limit_index; in SliceInDim()
1056 strides[dimno] = stride; in SliceInDim()
1160 XlaOp XlaBuilder::PadInDim(XlaOp operand, XlaOp padding_value, int64_t dimno, in PadInDim() argument
1165 auto* dims = padding_config.mutable_dimensions(dimno); in PadInDim()
4074 XlaOp PadInDim(XlaOp operand, XlaOp padding_value, int64_t dimno, in PadInDim() argument
4076 return operand.builder()->PadInDim(operand, padding_value, dimno, pad_lo, in PadInDim()
4118 int64_t stride, int64_t dimno) { in SliceInDim() argument
4120 stride, dimno); in SliceInDim()
|
| /external/tensorflow/tensorflow/compiler/xla/python/xla_extension/ |
| D | ops.pyi | 363 dimno: int) -> XlaOp: ...
|
| /external/tensorflow/tensorflow/compiler/xla/python/ |
| D | xla_client_test.py | 1382 dimno=1)
1389 dimno=0)
|