Merge pull request #30380 from trevor-m:tmorris_flatten
PiperOrigin-RevId: 261773167
This commit is contained in:
commit
0e8d6b4bea
@ -580,9 +580,15 @@ class Flatten(Layer):
|
|||||||
permutation.append(1)
|
permutation.append(1)
|
||||||
inputs = array_ops.transpose(inputs, perm=permutation)
|
inputs = array_ops.transpose(inputs, perm=permutation)
|
||||||
|
|
||||||
outputs = array_ops.reshape(
|
input_shape = inputs.shape
|
||||||
inputs, (tensor_shape.dimension_value(inputs.shape[0]) or
|
if input_shape[1:].is_fully_defined():
|
||||||
array_ops.shape(inputs)[0], -1))
|
flattened_dim = tensor_shape.dimension_value(
|
||||||
|
np.prod(input_shape[1:], dtype=int))
|
||||||
|
outputs = array_ops.reshape(inputs, (-1, flattened_dim))
|
||||||
|
else:
|
||||||
|
outputs = array_ops.reshape(
|
||||||
|
inputs, (tensor_shape.dimension_value(inputs.shape[0]) or
|
||||||
|
array_ops.shape(inputs)[0], -1))
|
||||||
if not context.executing_eagerly():
|
if not context.executing_eagerly():
|
||||||
outputs.set_shape(self.compute_output_shape(inputs.shape))
|
outputs.set_shape(self.compute_output_shape(inputs.shape))
|
||||||
return outputs
|
return outputs
|
||||||
|
Loading…
x
Reference in New Issue
Block a user