Function _convert_flatten_batch

tensorflow/python/ops/parallel_for/pfor.py:1501–1510

(pfor_input, op_type, dims)

@RegisterPForWithArgs("MaxPoolGradGrad", dims=[0, 1, 2])
@RegisterPForWithArgs("SoftmaxCrossEntropyWithLogits", dims=[0, 1])
def _convert_flatten_batch(pfor_input, op_type, dims):
  del op_type
  inputs = _inputs_with_flattening(pfor_input, dims)
  outputs = _create_op(
      pfor_input.op_type,
      inputs, [x.dtype for x in pfor_input.outputs],
      attrs=pfor_input.op.node_def.attr).outputs
  n = pfor_input.pfor.loop_len_vector
  outputs = [_unflatten_first_dim(x, n) for x in outputs]
  return [wrap(x, True) for x in outputs]


_channel_flatten_input_cache = {}