Searched refs:num_index_dims (Results 1 – 3 of 3) sorted by relevance
46 int64 num_index_dims; in XlaGather() local50 num_index_dims = indices_shape.dim_size(indices_shape.dims() - 1); in XlaGather()55 num_index_dims = 1; in XlaGather()66 input_shape_post_axis.RemoveDimRange(0, axis + num_index_dims); in XlaGather()83 for (int64 i = 0; i < num_index_dims; ++i) { in XlaGather()121 if (axis <= i && i < (axis + num_index_dims)) { in XlaGather()132 } else if (i >= (axis + num_index_dims)) { in XlaGather()135 dim_numbers.add_offset_dims(i + indices_rank - num_index_dims); in XlaGather()141 for (int64 i = axis; i < axis + num_index_dims; i++) { in XlaGather()212 const int64 num_index_dims = in Compile() local[all …]
41 const int64 num_index_dims = indices_shape.dim_size(indices_shape.dims() - 1); in ValidateUpdateShape() local51 ", num_index_dims: ", num_index_dims, ", and batch_dim: ", batch_dim); in ValidateUpdateShape()56 num_index_dims + (updates_shape.dims() - batch_dim)) { in ValidateUpdateShape()60 batch_dim + buffer_shape.dims() - num_index_dims) { in ValidateUpdateShape()70 buffer_shape.dim_size(d + num_index_dims)) { in ValidateUpdateShape()
47 int64 num_index_dims = 1; in XlaScatter() local50 num_index_dims = indices_dims.back(); in XlaScatter()51 if (num_index_dims > buffer_shape.rank()) { in XlaScatter()73 for (int64 i = 0; i < num_index_dims; ++i) { in XlaScatter()145 int64 num_window_dims_in_updates = buffer_rank - num_index_dims; in XlaScatter()152 for (int64 dim = num_index_dims; dim < buffer_rank; ++dim) { in XlaScatter()169 for (int64 i = 0; i < num_index_dims; ++i) { in XlaScatter()