Change np.asarray to np.stack.
It should be equivalent but in practice this speeds up things due to not reallocating memory.
@@ -83,6 +83,12 @@ def stack_sequence_fields(sequence: Iterable[ElementType]) -> ElementType:
if not sequence:
raise ValueError('Input sequence must not be empty')
+ # Default to asarray when arrays don't have the same shape to be compatible
+ # with old behaviour.
+ # TODO(b/169306678) make this more elegant.
+ try:
+ return fast_map_structure(lambda *values: np.stack(values), *sequence)
+ except ValueError:
return fast_map_structure(lambda *values: np.asarray(values), *sequence)
