Searched refs:new_batch_dim (Results 1 – 2 of 2) sorted by relevance
138 new_batch_dim = self._compute_static_batch_dim()142 lambda ts: ts._unbatch()._batch(new_batch_dim),170 new_batch_dim = tensor_util.constant_value(self._batch_sizes)171 if new_batch_dim is None:174 if isinstance(new_batch_dim, np.ndarray):175 if len(new_batch_dim.shape) == 1:176 if np.all(new_batch_dim == new_batch_dim[0]):177 new_batch_dim = new_batch_dim[0]180 elif len(new_batch_dim.shape) > 1:183 if self._may_form_partial_batches(new_batch_dim):[all …]
152 HloInstruction* select_val, int64 new_batch_dim, int64 new_space_dim,488 int64 new_batch_dim, new_spatial_dim; in BringSpaceNextToBatch() local498 new_batch_dim = pushed_counter; in BringSpaceNextToBatch()516 activations_batch_dim = new_batch_dim; in BringSpaceNextToBatch()530 new_batch_dim = pushed_counter; in BringSpaceNextToBatch()550 activations_batch_dim = new_batch_dim; in BringSpaceNextToBatch()831 const int64 new_batch_dim = in CanPropagate() local835 const int64 lhs_batch = first_operand->shape().dimensions(new_batch_dim); in CanPropagate()865 const int64 new_batch_dim = in CanPropagate() local867 const int64 rhs_batch = second_operand->shape().dimensions(new_batch_dim); in CanPropagate()[all …]