Returns a structure with `x[index]` for each tensor `x` in the structure. Args: struct: A structure of tensors. index: A scalar integer tensor. Performance is better if `index` is on the host memory. Returns: A structure of tensors congruent to `struct`. For each key in `
(struct, index)
| 66 | |
| 67 | |
| 68 | def _Index(struct, index): |
| 69 | """Returns a structure with `x[index]` for each tensor `x` in the structure. |
| 70 | |
| 71 | Args: |
| 72 | struct: A structure of tensors. |
| 73 | index: A scalar integer tensor. Performance is better if `index` is |
| 74 | on the host memory. |
| 75 | |
| 76 | Returns: |
| 77 | A structure of tensors congruent to `struct`. |
| 78 | For each key in `ret`, `rets[key] = struct[key][index]`. |
| 79 | """ |
| 80 | index = ops.convert_to_tensor(index) |
| 81 | index.get_shape().assert_has_rank(0) |
| 82 | return nest.map_structure(lambda x: array_ops.gather(x, index), struct) |
| 83 | |
| 84 | |
| 85 | def _Update(struct_acc, struct_x, t): |
no test coverage detected