Method max

experiments/olmo/preprocessing/preprocessor_utils.py:79–90

(*other: 'TensorSpec')

    @staticmethod
    def max(*other: 'TensorSpec'):
        other = [o for o in other if o is not None]
        if not other:
            return None
        rank = len(other[0].shape)
        dtype = other[0].dtype
        assert all(rank == len(o.shape) for o in other)
        assert all(dtype == o.dtype for o in other)
        return TensorSpec(
            [max([o.shape[i] for o in other]) for i in range(rank)],
            dtype
        )

    def max_dictionaries(*other: Dict[str, 'TensorSpec']):
        return {k: TensorSpec.max(*(o[k] for o in other if k in o)) for k in get_all_keys(other)}