(pfor_input)
| 1952 | @RegisterPFor("Gather") |
| 1953 | @RegisterPFor("GatherV2") |
| 1954 | def _convert_gather(pfor_input): |
| 1955 | param, param_stacked, _ = pfor_input.input(0) |
| 1956 | indices, indices_stacked, _ = pfor_input.input(1) |
| 1957 | op_type = pfor_input.op_type |
| 1958 | if op_type == "Gather": |
| 1959 | validate_indices = pfor_input.get_attr("validate_indices") |
| 1960 | axis = 0 |
| 1961 | else: |
| 1962 | validate_indices = None |
| 1963 | axis = pfor_input.unstacked_input(2) |
| 1964 | axis_value = tensor_util.constant_value(axis) |
| 1965 | if axis_value is not None: |
| 1966 | axis = axis_value |
| 1967 | if indices_stacked and not param_stacked: |
| 1968 | if indices is pfor_input.pfor.all_indices and axis == 0: |
| 1969 | param_shape0 = param.shape.dims[0].value |
| 1970 | indices_shape0 = indices.shape.dims[0].value |
| 1971 | if param_shape0 is not None and indices_shape0 == param_shape0: |
| 1972 | # Note that with loops and conditionals, indices may not be contiguous. |
| 1973 | # However they will be sorted and unique. So if the shape matches, then |
| 1974 | # it must be picking up all the rows of param. |
| 1975 | return wrap(param, True) |
| 1976 | # TODO(agarwal): use array_ops.slice here. |
| 1977 | output = array_ops.gather( |
| 1978 | param, indices, validate_indices=validate_indices, axis=axis) |
| 1979 | if axis != 0: |
| 1980 | axis = control_flow_ops.cond( |
| 1981 | axis < 0, lambda: axis + array_ops.rank(param), lambda: axis) |
| 1982 | order = array_ops.concat( |
| 1983 | [[axis], |
| 1984 | math_ops.range(axis), |
| 1985 | math_ops.range(axis + 1, array_ops.rank(output))], |
| 1986 | axis=0) |
| 1987 | output = control_flow_ops.cond( |
| 1988 | math_ops.equal(axis, 0), lambda: output, |
| 1989 | lambda: array_ops.transpose(output, order)) |
| 1990 | return wrap(output, True) |
| 1991 | if param_stacked: |
| 1992 | loop_len_vector = pfor_input.pfor.loop_len_vector |
| 1993 | pfor_input.stack_inputs(stack_indices=[1]) |
| 1994 | indices = pfor_input.stacked_input(1) |
| 1995 | param_flat = _flatten_first_two_dims(param) |
| 1996 | |
| 1997 | # Recompute indices to handle stacked param. |
| 1998 | indices_offset = math_ops.range( |
| 1999 | loop_len_vector[0]) * array_ops.shape(param)[1] |
| 2000 | # Reshape indices_offset to allow broadcast addition |
| 2001 | ones = array_ops.ones([array_ops.rank(indices) - 1], dtype=dtypes.int32) |
| 2002 | new_shape = array_ops.concat([loop_len_vector, ones], axis=0) |
| 2003 | indices_offset = array_ops.reshape(indices_offset, new_shape) |
| 2004 | indices += indices_offset |
| 2005 | |
| 2006 | # TODO(agarwal): handle axis != 0. May need to transpose param or |
| 2007 | # array_ops.gather_nd. |
| 2008 | if isinstance(axis, ops.Tensor): |
| 2009 | axis_value = tensor_util.constant_value(axis) |
| 2010 | else: |
| 2011 | try: |
nothing calls this directly
no test coverage detected