Repository: activeloopai/Hub
File: deeplake/core/chunk_engine.py (Python)

Code before:
        flatten: bool = False,
    ):
        index_0, sample_index = index.values[0], index.values[1:]
        sample_indices = list(
            index_0.indices(self._sequence_length or self.num_samples)
        )
        num_samples = len(sample_indices)

        sample_ndim = self.ndim() - 1

        for i, idx in enumerate(sample_indices):
            if self.tensor_meta.htype in ("text", "json"):
                shape = (1,)
            else:
                if sample_shape_provider:
                    try:
                        shape = sample_shape_provider(idx)  # type: ignore
                    except (
                        IndexError
                    ):  # Happens during transforms, sample shape tensor is not populated yet
                        shape = self.read_shape_for_sample(idx)  # type: ignore

                    if isinstance(shape, tuple) and shape == ():
                        shape = (0,)
                    if self.is_sequence:
                        if not flatten:
                            if sample_index and not sample_index[0].subscriptable():
                                shape = (1, *tuple(shape[sample_index[0].value].tolist()))  # type: ignore
                            else:
                                is_same = np.all(shape == shape[0, :], axis=0)  # type: ignore
                                shape = (len(shape),) + (
                                    tuple(
                                        int(shape[0, i])  # type: ignore
                                        if is_same[i]  # type: ignore
                                        else -1
                                        for i in range(shape.shape[1])  # type: ignore
                                    )
                                    or (1,)
                                )
                else:
                    self.check_link_ready()
                    shape = self.read_shape_for_sample(idx)  # type: ignore
                    # if link verification was not done
                    if len(shape) > sample_ndim:
                        sample_ndim = len(shape)
                        sample_shapes = np.zeros(
                            (num_samples, sample_ndim), dtype=np.int32
                        )
                if flatten:
                    start, end = self.sequence_encoder[i]
                    sample_shapes[start:end] = shape

Review comment (suggestion):
can we combine 
```suggestion
                        if self.is_sequence and not flatten:
```

Code after:
        except (
            IndexError
        ):  # Happens during transforms, sample shape tensor is not populated yet
            shape = self.read_shape_for_sample(idx)  # type: ignore

        if isinstance(shape, tuple) and shape == ():
            shape = (0,)
        if self.is_sequence and not flatten:
            shape = self._merge_seq_shape(shape, sample_index)
        return shape

    def _merge_seq_shape(self, shape, sample_index):
        """Merges shapes of sequence items into one shape"""
        if sample_index and not sample_index[0].subscriptable():
            shape = (1, *tuple(shape[sample_index[0].value].tolist()))  # type: ignore
        else:
            is_same = np.all(shape == shape[0, :], axis=0)  # type: ignore
            shape = (len(shape),) + (
                tuple(
                    int(shape[0, i])  # type: ignore
                    if is_same[i]  # type: ignore
                    else -1
                    for i in range(shape.shape[1])  # type: ignore
                )
                or (1,)
            )
        return shape

    def _populate_sample_shapes(
        self,
        sample_shapes: np.ndarray,
        index: Index,
        sample_shape_provider: Optional[Callable] = None,
        flatten: bool = False,
    ):
        index_0, sample_index = index.values[0], index.values[1:]
        sample_indices = list(
            index_0.indices(self._sequence_length or self.num_samples)
        )
        num_samples = len(sample_indices)

        sample_ndim = self.ndim() - 1

        for i, idx in enumerate(sample_indices):
            if self.tensor_meta.htype in ("text", "json"):
                shape = (1,)
            elif sample_shape_provider:
                shape = self._get_sample_shape_from_provider(
                    sample_shape_provider, idx, sample_index, flatten
                )
            else: