Searched refs:input_partition_dims (Results 1 – 4 of 4) sorted by relevance
123 input_partition_dims=None, argument134 if input_partition_dims is not None:135 if len(input_partition_dims) != 1 and len(input_partition_dims) != 2:184 input_partition_dims=input_partition_dims,
911 input_partition_dims = inputs_structure_recorder.flattened_input_dims913 input_partition_dims += [None] * len(signals)918 input_partition_dims=input_partition_dims,1072 def __init__(self, input_partition_dims=None): argument1077 if input_partition_dims:1079 assert len(input_partition_dims) <= 2, 'must have 1 or 2 elements.'1080 if len(input_partition_dims) == 2:1081 self._feature_dims, self._label_dims = input_partition_dims1083 self._feature_dims = input_partition_dims[0]1204 ctx.input_partition_dims)
748 input_partition_dims=None, argument758 self._input_partition_dims = input_partition_dims
334 def input_partition_dims(self): member in _InternalTPUContext335 return self._config.tpu_config.input_partition_dims