dask-array 0.1.0__py3-none-any.whl

dask_array/tests/test_slice_through_overlap.py

@@ -0,0 +1,366 @@
+"""Tests for slice pushdown through MapOverlap.
+
+These tests verify that slicing operations can be pushed through map_overlap
+operations, reducing computation by slicing input arrays before applying
+overlap boundaries.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import dask_array as da
+from dask_array._test_utils import assert_eq
+
+
+def add_neighbors(x):
+    """Add neighboring values along axis 0. Uses overlap data."""
+    result = x.copy()
+    if x.shape[0] > 2:
+        result[1:-1] = x[:-2] + x[1:-1] + x[2:]
+    return result
+
+
+# =============================================================================
+# Case 1: Slice on non-overlap axis (should push through)
+# =============================================================================
+
+
+def test_slice_through_overlap_non_overlap_axis():
+    """Slice on axis without overlap pushes through."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    # Overlap only on axis 0
+    result = x.map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    # Slice on axis 1 (no overlap) - should be equivalent to slicing input first
+    sliced = result[:, :20]
+    expected = x[:, :20].map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    # Verify expression structure matches
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+def test_slice_through_overlap_middle_slice():
+    """Slice in the middle of non-overlap axis."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    # Middle slice on axis 1 (no overlap)
+    sliced = result[:, 30:70]
+    expected = x[:, 30:70].map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+def test_slice_through_overlap_correctness():
+    """Verify slice through overlap produces correct values."""
+    arr = np.arange(64).reshape((8, 8)).astype(float)
+    x = da.from_array(arr, chunks=(4, 4))
+
+    result = x.map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    # Slice on axis 1
+    sliced = result[:, 2:6]
+    expected = x[:, 2:6].map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+# =============================================================================
+# Case 2: Slice on overlap axis (pushes through with padding)
+# =============================================================================
+
+
+def test_slice_on_overlap_axis_pushes_with_padding():
+    """Slice on axis with overlap pushes through with padded input."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    # Slice on axis 0 (has overlap) - should push through with padded input
+    # [:50] with depth=2 needs input [:52], then trim to [:50]
+    sliced = result[:50, :]
+    expected = x[:52, :].map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")[:50, :]
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+def test_slice_on_both_axes_one_has_overlap():
+    """Slice on both axes when one has overlap."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+    sliced = result[:50, :50]
+
+    # Axis 1 has no overlap: slice pushes directly
+    # Axis 0 has depth=2: need padded input [:52], then trim to [:50]
+    expected = x[:52, :50].map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")[:50, :]
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+# =============================================================================
+# Case 3: Multi-dimensional overlap
+# =============================================================================
+
+
+def add_neighbors_2d(x):
+    """Add neighboring values along both axes. Uses overlap data."""
+    result = x.copy()
+    if x.shape[0] > 2:
+        result[1:-1, :] += x[:-2, :] + x[2:, :]
+    if x.shape[1] > 2:
+        result[:, 1:-1] += x[:, :-2] + x[:, 2:]
+    return result
+
+
+def test_slice_through_2d_overlap():
+    """Slice through 2D overlap - pushes when beneficial."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.map_overlap(add_neighbors_2d, depth={0: 1, 1: 1}, boundary="none")
+
+    # Slice on axis 1 with depth=1 needs input [:, :41], then trim to [:, :40]
+    sliced = result[:, :40]
+    expected = x[:, :41].map_overlap(add_neighbors_2d, depth={0: 1, 1: 1}, boundary="none")[:, :40]
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+def test_slice_through_2d_overlap_middle():
+    """Middle slice through 2D overlap on non-overlap dimension."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    # Overlap only on axis 0
+    result = x.map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    # Middle slice on axis 1 (no overlap)
+    sliced = result[:, 25:75]
+    expected = x[:, 25:75].map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+def test_slice_through_1d_overlap_on_3d_array():
+    """Slice on multiple non-overlap axes."""
+    arr = np.arange(1000).reshape((10, 10, 10)).astype(float)
+    x = da.from_array(arr, chunks=(5, 5, 5))
+
+    # Overlap only on axis 0
+    result = x.map_overlap(add_neighbors, depth={0: 1, 1: 0, 2: 0}, boundary="none")
+
+    # Slice on axes 1 and 2 (neither has overlap)
+    sliced = result[:, :3, :3]
+    expected = x[:, :3, :3].map_overlap(add_neighbors, depth={0: 1, 1: 0, 2: 0}, boundary="none")
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+# =============================================================================
+# Case 4: Asymmetric overlap
+# =============================================================================
+
+
+def test_slice_through_asymmetric_overlap():
+    """Slice through asymmetric overlap (different left/right depth)."""
+    arr = np.arange(64).reshape((8, 8)).astype(float)
+    x = da.from_array(arr, chunks=(4, 4))
+
+    # Asymmetric overlap on axis 0
+    result = x.map_overlap(add_neighbors, depth={0: (2, 1), 1: 0}, boundary="none")
+
+    # Slice on axis 1 (no overlap)
+    sliced = result[:, 2:6]
+    expected = x[:, 2:6].map_overlap(add_neighbors, depth={0: (2, 1), 1: 0}, boundary="none")
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+def test_slice_on_asymmetric_overlap_axis_pushes():
+    """Slice on axis with asymmetric overlap pushes through with padding."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.map_overlap(add_neighbors, depth={0: (2, 1), 1: 0}, boundary="none")
+
+    # Slice axis 0 with asymmetric depth (2, 1) - needs extra 1 on right
+    # [:50] needs input [:51], then trim to [:50]
+    sliced = result[:50, :]
+    expected = x[:51, :].map_overlap(add_neighbors, depth={0: (2, 1), 1: 0}, boundary="none")[:50, :]
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+# =============================================================================
+# Case 5: Zero overlap (edge case)
+# =============================================================================
+
+
+def test_slice_through_zero_overlap():
+    """Slice through axis with zero overlap pushes through."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    # Zero overlap - no actual overlap computation needed
+    result = x.map_overlap(add_neighbors, depth=0, boundary="none")
+
+    # Slice on axis 0 - with zero overlap, slice should push through
+    sliced = result[:50, :]
+    expected = x[:50, :].map_overlap(add_neighbors, depth=0, boundary="none")
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+# =============================================================================
+# Case 6: Task reduction verification
+# =============================================================================
+
+
+def test_slice_through_overlap_reduces_tasks():
+    """Verify slice pushdown reduces task count."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    full = result
+    sliced = result[:, :10]  # Take only first 10 columns
+
+    full_tasks = len(full.optimize().__dask_graph__())
+    sliced_tasks = len(sliced.optimize().__dask_graph__())
+
+    # Sliced should have fewer tasks (processes 1 column of chunks vs 10)
+    assert sliced_tasks < full_tasks
+
+
+def test_slice_through_overlap_reduces_numblocks():
+    """Verify slice pushdown reduces number of output blocks."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+    sliced = result[:, :10]
+
+    # Full result: 10x10 chunks
+    assert result.numblocks == (10, 10)
+
+    # Sliced result: 10x1 chunks (only 1 column of blocks)
+    assert sliced.numblocks == (10, 1)
+
+
+# =============================================================================
+# Case 7: Correctness with computed values
+# =============================================================================
+
+
+@pytest.mark.parametrize(
+    "shape,chunks,depth,slice_",
+    [
+        # Start slices (:n form) on non-overlap axes
+        ((80, 80), (20, 20), {0: 2, 1: 0}, (slice(None), slice(20))),
+        ((80, 80), (20, 20), {0: 0, 1: 2}, (slice(20), slice(None))),
+        # Middle slices (k:n form) on non-overlap axes
+        ((80, 80), (20, 20), {0: 2, 1: 0}, (slice(None), slice(20, 60))),
+        ((80, 80), (20, 20), {0: 0, 1: 2}, (slice(20, 60), slice(None))),
+        # End slices (k: form) on non-overlap axes
+        ((80, 80), (20, 20), {0: 2, 1: 0}, (slice(None), slice(40, None))),
+        ((80, 80), (20, 20), {0: 0, 1: 2}, (slice(40, None), slice(None))),
+    ],
+)
+def test_slice_through_overlap_parametrized(shape, chunks, depth, slice_):
+    """Parametrized correctness tests for slice through overlap."""
+    arr = np.arange(np.prod(shape)).reshape(shape).astype(float)
+    x = da.from_array(arr, chunks=chunks)
+
+    result = x.map_overlap(add_neighbors, depth=depth, boundary="none")
+    sliced = result[slice_]
+
+    # Build expected: slice input first, then overlap
+    input_sliced = x[slice_]
+    expected = input_sliced.map_overlap(add_neighbors, depth=depth, boundary="none")
+
+    # Verify expression structure matches
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+# =============================================================================
+# Case 8: Special cases (trim=False, uniform depth)
+# =============================================================================
+
+
+def test_map_overlap_no_trim_slice_pushes():
+    """With trim=False, slice should push through to input."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    # With trim=False, there's no Trim wrapper, so slice can push through
+    result = x.map_overlap(add_neighbors, depth={0: 2}, boundary="none", trim=False)
+
+    # Slice on axis 1 (no overlap on axis 1) - pushes directly through
+    sliced = result[:, :30]
+    expected = x[:, :30].map_overlap(add_neighbors, depth={0: 2}, boundary="none", trim=False)
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+def test_map_overlap_uniform_depth_correctness():
+    """Test with uniform depth (int instead of dict).
+
+    When slicing on an axis with overlap, the optimization pads the input
+    slice to include data needed for overlap, then trims the output.
+    """
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.map_overlap(add_neighbors_2d, depth=2, boundary="none")
+    sliced = result[:, :30]
+
+    # Expected: pad input by depth on sliced axis, apply overlap, then trim
+    # [:, :30] with depth=2 needs input [:, :32] to preserve overlap semantics
+    expected = x[:, :32].map_overlap(add_neighbors_2d, depth=2, boundary="none")[:, :30]
+
+    assert sliced.expr.simplify()._name == expected.expr.simplify()._name
+
+
+# =============================================================================
+# Case 9: Value correctness verification
+# =============================================================================
+
+
+def test_slice_through_overlap_value_correctness():
+    """Verify optimized slice produces correct values."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.map_overlap(add_neighbors, depth={0: 2, 1: 0}, boundary="none")
+
+    # Slice on non-overlap axis
+    sliced = result[:, :50]
+
+    # Compare against unoptimized computation
+    full_result = result.compute()
+    assert_eq(sliced, full_result[:, :50])
+
+
+def test_slice_on_overlap_axis_value_correctness():
+    """Verify slice on overlap axis produces correct values."""
+    arr = np.arange(10000).reshape((100, 100)).astype(float)
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.map_overlap(add_neighbors_2d, depth=2, boundary="none")
+
+    # Slice on axis with overlap
+    sliced = result[:50, :50]
+
+    # Compare against unoptimized computation
+    full_result = result.compute()
+    assert_eq(sliced, full_result[:50, :50])

dask_array/tests/test_slice_through_reshape.py

@@ -0,0 +1,272 @@
+"""Tests for slice pushdown through Reshape expressions.
+
+Slice can push through Reshape when leading dimensions are preserved,
+i.e., the reshape only affects trailing dimensions.
+"""
+
+from __future__ import annotations
+
+import numpy as np
+import pytest
+
+import dask_array as da
+from dask_array._test_utils import assert_eq
+
+# =============================================================================
+# Case 1: Leading dimension preserved - slice should push through
+# =============================================================================
+
+
+def test_slice_through_reshape_leading_dim_preserved():
+    """Slice on preserved leading dimension pushes through reshape."""
+    arr = np.arange(60).reshape((10, 6))
+    x = da.from_array(arr, chunks=(5, 3))
+
+    # Reshape (10, 6) -> (10, 2, 3) preserves first dimension
+    result = x.reshape((10, 2, 3))[:3]
+    expected = x[:3].reshape((3, 2, 3))
+
+    # After simplification, both should have same structure
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+    assert_eq(result, arr.reshape((10, 2, 3))[:3])
+
+
+def test_slice_through_reshape_flatten_trailing():
+    """Slice on leading dimension when flattening trailing dims."""
+    arr = np.arange(60).reshape((10, 2, 3))
+    x = da.from_array(arr, chunks=(5, 2, 3))
+
+    # Reshape (10, 2, 3) -> (10, 6) preserves first dimension
+    result = x.reshape((10, 6))[:4]
+    expected = x[:4].reshape((4, 6))
+
+    # After optimization, both should produce equivalent graphs
+    # (Reshape vs ReshapeLowered difference resolved by lowering)
+    assert result.optimize()._name == expected.optimize()._name
+    assert_eq(result, arr.reshape((10, 6))[:4])
+
+
+def test_slice_through_reshape_multiple_leading_dims():
+    """Slice when multiple leading dimensions are preserved."""
+    arr = np.arange(120).reshape((4, 5, 6))
+    x = da.from_array(arr, chunks=(2, 5, 3))
+
+    # Reshape (4, 5, 6) -> (4, 5, 2, 3) preserves first two dimensions
+    result = x.reshape((4, 5, 2, 3))[:2, :3]
+    expected = x[:2, :3].reshape((2, 3, 2, 3))
+
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+    assert_eq(result, arr.reshape((4, 5, 2, 3))[:2, :3])
+
+
+def test_slice_through_reshape_middle_slice():
+    """Middle slice (start > 0) on preserved dimension."""
+    arr = np.arange(100).reshape((10, 10))
+    x = da.from_array(arr, chunks=(5, 5))
+
+    # Reshape (10, 10) -> (10, 2, 5) preserves first dimension
+    result = x.reshape((10, 2, 5))[3:7]
+    expected = x[3:7].reshape((4, 2, 5))
+
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+    assert_eq(result, arr.reshape((10, 2, 5))[3:7])
+
+
+# =============================================================================
+# Case 2: Dimension NOT preserved - slice should NOT push through
+# =============================================================================
+
+
+def test_slice_not_pushed_when_dim_changes():
+    """Slice blocked when the sliced dimension changes size."""
+    arr = np.arange(60).reshape((6, 10))
+    x = da.from_array(arr, chunks=(3, 5))
+
+    # Reshape (6, 10) -> (2, 3, 10) splits first dimension
+    # First dim changes from 6 to 2, so slice should NOT push through
+    result = x.reshape((2, 3, 10))[:1]
+
+    # Just verify correctness - can't push through
+    assert_eq(result, arr.reshape((2, 3, 10))[:1])
+
+
+def test_slice_not_pushed_through_flatten():
+    """Slice blocked when reshape completely flattens the array."""
+    arr = np.arange(100).reshape((10, 10))
+    x = da.from_array(arr, chunks=(5, 5))
+
+    # Reshape to 1D - no dimension correspondence
+    result = x.reshape((100,))[:30]
+
+    # Just verify correctness - can't push through
+    assert_eq(result, arr.reshape((100,))[:30])
+
+
+def test_slice_on_reshaped_axis_not_pushed():
+    """Slice on axis that was created by reshape doesn't push."""
+    arr = np.arange(60).reshape((10, 6))
+    x = da.from_array(arr, chunks=(5, 3))
+
+    # Reshape (10, 6) -> (10, 2, 3), then slice on new axis 1
+    result = x.reshape((10, 2, 3))[:, :1]
+
+    # Axis 1 is new (from splitting 6 -> 2, 3), can't push through
+    assert_eq(result, arr.reshape((10, 2, 3))[:, :1])
+
+
+# =============================================================================
+# Case 3: Slice with None (newaxis) - should still push through
+# =============================================================================
+
+
+def test_slice_with_none_pushes_through():
+    """Slice with None (newaxis) should push through and re-apply None."""
+    arr = np.arange(60).reshape((10, 6))
+    x = da.from_array(arr, chunks=(5, 3))
+
+    # Reshape (10, 6) -> (10, 2, 3), slice with None
+    result = x.reshape((10, 2, 3))[:5, None]
+    expected = x[:5].reshape((5, 2, 3))[:, None]
+
+    # Both should have same structure after simplification
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+    assert_eq(result, arr.reshape((10, 2, 3))[:5, None])
+
+
+def test_slice_with_none_at_end():
+    """Slice with None at end of index."""
+    arr = np.arange(60).reshape((10, 6))
+    x = da.from_array(arr, chunks=(5, 3))
+
+    result = x.reshape((10, 2, 3))[:5, :, :, None]
+    expected = x[:5].reshape((5, 2, 3))[:, :, :, None]
+
+    # Both should have same structure after simplification
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+    assert_eq(result, arr.reshape((10, 2, 3))[:5, :, :, None])
+
+
+def test_slice_with_multiple_nones():
+    """Slice with multiple Nones."""
+    arr = np.arange(60).reshape((10, 6))
+    x = da.from_array(arr, chunks=(5, 3))
+
+    result = x.reshape((10, 2, 3))[None, :5, None]
+    expected = x[:5].reshape((5, 2, 3))[None, :, None]
+
+    # Both should have same structure after simplification
+    assert result.expr.simplify()._name == expected.expr.simplify()._name
+    assert_eq(result, arr.reshape((10, 2, 3))[None, :5, None])
+
+
+def test_slice_with_none_correctness():
+    """Verify correctness with various None positions."""
+    arr = np.arange(120).reshape((10, 12))
+    x = da.from_array(arr, chunks=(5, 6))
+
+    # (10, 12) -> (10, 3, 4)
+    reshaped = x.reshape((10, 3, 4))
+
+    # Various slices with Nones
+    assert_eq(reshaped[:5, None, :, :], arr.reshape((10, 3, 4))[:5, None, :, :])
+    assert_eq(reshaped[None, :5], arr.reshape((10, 3, 4))[None, :5])
+    assert_eq(reshaped[:5, :, None, :], arr.reshape((10, 3, 4))[:5, :, None, :])
+
+
+# =============================================================================
+# Correctness tests with various shapes
+# =============================================================================
+
+
+@pytest.mark.parametrize(
+    "in_shape,out_shape,slice_",
+    [
+        # Trailing split
+        ((20, 6), (20, 2, 3), (slice(10),)),
+        ((20, 6), (20, 2, 3), (slice(5, 15),)),
+        ((20, 12), (20, 3, 4), (slice(None, 8),)),
+        # Trailing merge
+        ((20, 2, 3), (20, 6), (slice(10),)),
+        ((20, 4, 5), (20, 20), (slice(5, 15),)),
+        # Multiple preserved dims
+        ((10, 5, 6), (10, 5, 2, 3), (slice(5), slice(3))),
+        ((10, 5, 4), (10, 5, 2, 2), (slice(3, 8), slice(None, 4))),
+    ],
+)
+def test_slice_through_reshape_correctness(in_shape, out_shape, slice_):
+    """Parametrized correctness tests."""
+    arr = np.arange(np.prod(in_shape)).reshape(in_shape)
+    chunks = tuple(max(1, s // 2) for s in in_shape)
+    x = da.from_array(arr, chunks=chunks)
+
+    result = x.reshape(out_shape)[slice_]
+    expected = arr.reshape(out_shape)[slice_]
+
+    assert_eq(result, expected)
+
+
+# =============================================================================
+# Task reduction tests
+# =============================================================================
+
+
+def test_slice_through_reshape_reduces_tasks():
+    """Verify slice pushdown reduces task count."""
+    arr = np.ones((100, 10))
+    x = da.from_array(arr, chunks=(10, 5))
+
+    # Reshape preserves first dim, then slice
+    full = x.reshape((100, 2, 5))
+    sliced = x.reshape((100, 2, 5))[:10]
+
+    full_tasks = len(full.optimize().__dask_graph__())
+    sliced_tasks = len(sliced.optimize().__dask_graph__())
+
+    # Sliced should have fewer tasks (only 1/10 of chunks)
+    assert sliced_tasks < full_tasks
+
+
+def test_slice_through_reshape_reduces_numblocks():
+    """Verify slice pushdown reduces number of blocks."""
+    arr = np.ones((100, 20))
+    x = da.from_array(arr, chunks=(10, 10))
+
+    result = x.reshape((100, 4, 5))[:20]
+    optimized = result.optimize()
+
+    # Should only have 2 blocks in first dimension (20 / 10)
+    assert optimized.numblocks[0] == 2
+
+
+# =============================================================================
+# Expression structure tests
+# =============================================================================
+
+
+def test_expression_structure_slice_pushed():
+    """Verify slice is pushed through reshape in expression tree."""
+    from dask_array.slicing import SliceSlicesIntegers
+
+    x = da.ones((20, 6), chunks=(5, 3))
+    result = x.reshape((20, 2, 3))[:5]
+
+    # Before simplification: Slice(Reshape(...))
+    assert isinstance(result.expr, SliceSlicesIntegers)
+
+    # After simplification: slice should have pushed through
+    simplified = result.expr.simplify()
+
+    # Slice shouldn't be at root after pushdown
+    assert not isinstance(simplified, SliceSlicesIntegers)
+
+
+def test_expression_structure_slice_blocked():
+    """Verify slice is NOT pushed when dimension changes."""
+    x = da.ones((6, 10), chunks=(3, 5))
+
+    # Reshape (6, 10) -> (2, 3, 10) splits first dim
+    # First dim changes from 6 to 2, slice should not push through
+    result = x.reshape((2, 3, 10))[:1]
+
+    # Just verify correctness - the optimization doesn't apply here
+    assert_eq(result, np.ones((2, 3, 10))[:1])