Large memory usage of sum when using slices or the axis argument.

[aplund]

This may seem like a silly thing to do, but it's a somewhat minimal example of something that is confusing me. This script works for reducing the exceedingly large array, but fails on the sub-arrays:

import jax
import jax.numpy as jnp

import resource

resource.setrlimit(resource.RLIMIT_DATA, (4 * 1024**3, 8 * 1024**3))
resource.setrlimit(resource.RLIMIT_CORE, (0, 0))

jax.config.update("jax_enable_x64", True)


def tst_flat(n):
    big_array = jnp.mgrid.__getitem__((slice(10), ) * n)
    big_array_semiflat = jnp.reshape(big_array, (-1, 10))
    return jnp.sum(big_array_semiflat, dtype=jnp.int64)


def tst_slice(n, idx):
    big_array = jnp.mgrid.__getitem__((slice(10), ) * n)
    big_array_semiflat = jnp.reshape(big_array, (-1, 10))
    return jnp.sum(big_array_semiflat[..., idx], dtype=jnp.int64)


def tst_axis_arg(n):
    big_array = jnp.mgrid.__getitem__((slice(10), ) * n)
    big_array_semiflat = jnp.reshape(big_array, (-1, 10))
    return jnp.sum(big_array_semiflat, axis=0, dtype=jnp.int64)


tst_flat_jit = jax.jit(tst_flat, static_argnums=0)
tst_slice_jit = jax.jit(tst_slice, static_argnums=(0, 1))
tst_axis_arg_jit = jax.jit(tst_axis_arg, static_argnums=0)

m = 10
result_flat = tst_flat_jit(m)
result_slice = [tst_slice_jit(m, k) for k in range(m+1)]
result_axis_arg = tst_axis_arg_jit(m)

print(f"flat count: {result_flat}")
print(f"slice count: {result_slice}")
print(f"axis_arg count: {result_axis_arg}")

The output I get looks like this:

flat count: 40500000000
---------------------------------------------------------------------------
XlaRuntimeError                           Traceback (most recent call last)
File ~/Work/00Projects/Xanadu/StateGen/GaussianBargmannFock/jax_test1.py:40
     37 result_axis_arg = tst_axis_arg_jit(m)
     39 print(f"flat count: {result_flat}")
---> 40 print(f"slice count: {result_slice}")
     41 print(f"axis_arg count: {result_axis_arg}")

File ~/.pyenv/versions/3.12.4/envs/mrmustard/lib/python3.12/site-packages/jax/_src/array.py:401, in ArrayImpl.__repr__(self)
    399   s = f"[], shape={self.shape}"
    400 else:
--> 401   s = np.array2string(self._value, prefix=prefix, suffix=',',
    402                       separator=', ', max_line_width=line_width)
    403 last_line_len = len(s) - s.rfind('\n') + 1
    404 sep = ' '

File ~/.pyenv/versions/3.12.4/envs/mrmustard/lib/python3.12/site-packages/jax/_src/profiler.py:333, in annotate_function.<locals>.wrapper(*args, **kwargs)
    330 @wraps(func)
    331 def wrapper(*args, **kwargs):
    332   with TraceAnnotation(name, **decorator_kwargs):
--> 333     return func(*args, **kwargs)
    334   return wrapper

File ~/.pyenv/versions/3.12.4/envs/mrmustard/lib/python3.12/site-packages/jax/_src/array.py:627, in ArrayImpl._value(self)
    625 if self._npy_value is None:
    626   if self.is_fully_replicated:
--> 627     self._npy_value = self._single_device_array_to_np_array()
    628     self._npy_value.flags.writeable = False
    629     return cast(np.ndarray, self._npy_value)

XlaRuntimeError: FAILED_PRECONDITION: Buffer Definition Event: Error preparing computation: %sOut of memory allocating 7425000000 bytes.

Surely, doing the reduction with fewer elements would take less memory rather than more (I would have thought).

[jakevdp]

It looks like in the first call you're allocating a very large array ($10^{11}$ floating point elements consumes about 400 GB) and then running into a memory error when allocating additional arrays.

Perhaps you need to ensure that the very large buffer is deleted before attempting further computations?

[aplund]

Perhaps you need to ensure that the very large buffer is deleted before attempting further computations?

I'm somewhat confused as to how this might happen. Do you mean, the local variables of the jitted functions remain allocated? One must explicitly 'del' them?

[jakevdp]

Ah, right, I missed the fact that this was jit-compiled. In that case there shouldn't be any intermediate variables floating around.

To see what's going on, I generated the compiled HLO for your functions. Here's the flat version:

print(tst_flat_jit.lower(10).compile().as_text())

Output

HloModule jit_tst_flat, is_scheduled=true, entry_computation_layout={()->s64[]}, allow_spmd_sharding_propagation_to_output={true}

%region_0.25 (Arg_0.26: s64[], Arg_1.27: s64[]) -> s64[] {
  %Arg_0.26 = s64[] parameter(0), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum"}
  %Arg_1.27 = s64[] parameter(1), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum"}
  ROOT %add.28 = s64[] add(s64[] %Arg_0.26, s64[] %Arg_1.27), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
}

%fused_computation () -> s64[] {
  %iota.30 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=1, metadata={op_name="jit(tst_flat)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=13}
  %constant.0 = s64[] constant(0)
  %reduce.20 = s64[] reduce(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.30, s64[] %constant.0), dimensions={0,1,2,3,4,5,6,7,8,9,10}, to_apply=%region_0.25, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %iota.29 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=2, metadata={op_name="jit(tst_flat)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=13}
  %reduce.19 = s64[] reduce(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.29, s64[] %constant.0), dimensions={0,1,2,3,4,5,6,7,8,9,10}, to_apply=%region_0.25, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %add.17 = s64[] add(s64[] %reduce.20, s64[] %reduce.19), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %iota.28 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=3, metadata={op_name="jit(tst_flat)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=13}
  %reduce.18 = s64[] reduce(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.28, s64[] %constant.0), dimensions={0,1,2,3,4,5,6,7,8,9,10}, to_apply=%region_0.25, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %add.16 = s64[] add(s64[] %add.17, s64[] %reduce.18), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %iota.27 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=4, metadata={op_name="jit(tst_flat)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=13}
  %reduce.17 = s64[] reduce(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.27, s64[] %constant.0), dimensions={0,1,2,3,4,5,6,7,8,9,10}, to_apply=%region_0.25, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %add.15 = s64[] add(s64[] %add.16, s64[] %reduce.17), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %iota.26 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=5, metadata={op_name="jit(tst_flat)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=13}
  %reduce.16 = s64[] reduce(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.26, s64[] %constant.0), dimensions={0,1,2,3,4,5,6,7,8,9,10}, to_apply=%region_0.25, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %add.14 = s64[] add(s64[] %add.15, s64[] %reduce.16), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %iota.25 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=6, metadata={op_name="jit(tst_flat)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=13}
  %reduce.15 = s64[] reduce(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.25, s64[] %constant.0), dimensions={0,1,2,3,4,5,6,7,8,9,10}, to_apply=%region_0.25, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %add.13 = s64[] add(s64[] %add.14, s64[] %reduce.15), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %iota.24 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=7, metadata={op_name="jit(tst_flat)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=13}
  %reduce.14 = s64[] reduce(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.24, s64[] %constant.0), dimensions={0,1,2,3,4,5,6,7,8,9,10}, to_apply=%region_0.25, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %add.12 = s64[] add(s64[] %add.13, s64[] %reduce.14), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %iota.23 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=8, metadata={op_name="jit(tst_flat)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=13}
  %reduce.13 = s64[] reduce(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.23, s64[] %constant.0), dimensions={0,1,2,3,4,5,6,7,8,9,10}, to_apply=%region_0.25, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %add.11 = s64[] add(s64[] %add.12, s64[] %reduce.13), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %iota.22 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=9, metadata={op_name="jit(tst_flat)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=13}
  %reduce.12 = s64[] reduce(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.22, s64[] %constant.0), dimensions={0,1,2,3,4,5,6,7,8,9,10}, to_apply=%region_0.25, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %add.10 = s64[] add(s64[] %add.11, s64[] %reduce.12), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  %iota.21 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=10, metadata={op_name="jit(tst_flat)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=13}
  %reduce.11 = s64[] reduce(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.21, s64[] %constant.0), dimensions={0,1,2,3,4,5,6,7,8,9,10}, to_apply=%region_0.25, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
  ROOT %add.9 = s64[] add(s64[] %add.10, s64[] %reduce.11), metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
}

ENTRY %main.30 () -> s64[] {
  ROOT %fusion = s64[] fusion(), kind=kLoop, calls=%fused_computation, metadata={op_name="jit(tst_flat)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=15}
}

And here's the slice version:

print(tst_slice_jit.lower(10, 0).compile().as_text())

Output

HloModule jit_tst_slice, is_scheduled=true, entry_computation_layout={()->s64[]}, allow_spmd_sharding_propagation_to_output={true}

%region_0.27 (Arg_0.28: s64[], Arg_1.29: s64[]) -> s64[] {
  %Arg_0.28 = s64[] parameter(0), metadata={op_name="jit(tst_slice)/jit(main)/reduce_sum"}
  %Arg_1.29 = s64[] parameter(1), metadata={op_name="jit(tst_slice)/jit(main)/reduce_sum"}
  ROOT %add.30 = s64[] add(s64[] %Arg_0.28, s64[] %Arg_1.29), metadata={op_name="jit(tst_slice)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=21}
}

%parallel_reduce-window (p: s64[10000000000,1], p.1: s64[]) -> s64[312500000,1] {
  %p = s64[10000000000,1]{1,0} parameter(0)
  %p.1 = s64[] parameter(1)
  ROOT %reduce-window.clone = s64[312500000,1]{1,0} reduce-window(s64[10000000000,1]{1,0} %p, s64[] %p.1), window={size=32x1 stride=32x1}, to_apply=%region_0.27, backend_config={"outer_dimension_partitions":["2"]}
}

%parallel_reduce-window.1 (p.2: s64[312500000,1], p.3: s64[]) -> s64[9765625,1] {
  %p.2 = s64[312500000,1]{1,0} parameter(0)
  %p.3 = s64[] parameter(1)
  ROOT %reduce-window.1.clone = s64[9765625,1]{1,0} reduce-window(s64[312500000,1]{1,0} %p.2, s64[] %p.3), window={size=32x1 stride=32x1}, to_apply=%region_0.27, backend_config={"outer_dimension_partitions":["2"]}
}

%parallel_reduce-window.2 (p.4: s64[9765625,1], p.5: s64[]) -> s64[305176,1] {
  %p.4 = s64[9765625,1]{1,0} parameter(0)
  %p.5 = s64[] parameter(1)
  ROOT %reduce-window.2.clone = s64[305176,1]{1,0} reduce-window(s64[9765625,1]{1,0} %p.4, s64[] %p.5), window={size=32x1 stride=32x1 pad=3_4x0_0}, to_apply=%region_0.27, backend_config={"outer_dimension_partitions":["2"]}
}

%parallel_reduce-window.3 (p.6: s64[305176,1], p.7: s64[]) -> s64[9537,1] {
  %p.6 = s64[305176,1]{1,0} parameter(0)
  %p.7 = s64[] parameter(1)
  ROOT %reduce-window.3.clone = s64[9537,1]{1,0} reduce-window(s64[305176,1]{1,0} %p.6, s64[] %p.7), window={size=32x1 stride=32x1 pad=4_4x0_0}, to_apply=%region_0.27, backend_config={"outer_dimension_partitions":["2"]}
}

%fused_computation.clone () -> s64[10000000000,1] {
  %iota.40 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=1, metadata={op_name="jit(tst_slice)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %iota.39 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=2, metadata={op_name="jit(tst_slice)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %iota.38 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=3, metadata={op_name="jit(tst_slice)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %iota.37 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=4, metadata={op_name="jit(tst_slice)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %iota.36 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=5, metadata={op_name="jit(tst_slice)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %iota.35 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=6, metadata={op_name="jit(tst_slice)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %iota.34 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=7, metadata={op_name="jit(tst_slice)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %iota.33 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=8, metadata={op_name="jit(tst_slice)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %iota.32 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=9, metadata={op_name="jit(tst_slice)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %iota.31 = s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} iota(), iota_dimension=10, metadata={op_name="jit(tst_slice)/jit(main)/broadcast_in_dim" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %concatenate.1 = s64[10,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} concatenate(s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.40, s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.39, s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.38, s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.37, s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.36, /*index=5*/s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.35, s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.34, s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.33, s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.32, s64[1,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %iota.31), dimensions={0}, metadata={op_name="jit(tst_slice)/jit(main)/concatenate" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  %bitcast.2 = s64[10000000000,10]{1,0} bitcast(s64[10,10,10,10,10,10,10,10,10,10,10]{10,9,8,7,6,5,4,3,2,1,0} %concatenate.1), metadata={op_name="jit(tst_slice)/jit(main)/concatenate" source_file="<ipython-input-1-5fe284a4623f>" source_line=19}
  ROOT %slice.1 = s64[10000000000,1]{1,0} slice(s64[10000000000,10]{1,0} %bitcast.2), slice={[0:10000000000], [0:1]}, metadata={op_name="jit(tst_slice)/jit(main)/slice" source_file="<ipython-input-1-5fe284a4623f>" source_line=21}
}

%parallel_fusion () -> s64[10000000000,1] {
  ROOT %fusion.clone = s64[10000000000,1]{1,0} fusion(), kind=kLoop, calls=%fused_computation.clone, metadata={op_name="jit(tst_slice)/jit(main)/slice" source_file="<ipython-input-1-5fe284a4623f>" source_line=21}, backend_config={"outer_dimension_partitions":["2"]}
}

ENTRY %main.32 () -> s64[] {
  %call.4 = s64[10000000000,1]{1,0} call(), to_apply=%parallel_fusion
  %constant.1 = s64[] constant(0)
  %call = s64[312500000,1]{1,0} call(s64[10000000000,1]{1,0} %call.4, s64[] %constant.1), to_apply=%parallel_reduce-window
  %call.1 = s64[9765625,1]{1,0} call(s64[312500000,1]{1,0} %call, s64[] %constant.1), to_apply=%parallel_reduce-window.1
  %call.2 = s64[305176,1]{1,0} call(s64[9765625,1]{1,0} %call.1, s64[] %constant.1), to_apply=%parallel_reduce-window.2
  %call.3 = s64[9537,1]{1,0} call(s64[305176,1]{1,0} %call.2, s64[] %constant.1), to_apply=%parallel_reduce-window.3
  %reduce-window.4 = s64[299,1]{1,0} reduce-window(s64[9537,1]{1,0} %call.3, s64[] %constant.1), window={size=32x1 stride=32x1 pad=15_16x0_0}, to_apply=%region_0.27
  %reduce-window.5 = s64[10,1]{1,0} reduce-window(s64[299,1]{1,0} %reduce-window.4, s64[] %constant.1), window={size=32x1 stride=32x1 pad=10_11x0_0}, to_apply=%region_0.27
  ROOT %reduce.31 = s64[] reduce(s64[10,1]{1,0} %reduce-window.5, s64[] %constant.1), dimensions={0,1}, to_apply=%region_0.27, metadata={op_name="jit(tst_slice)/jit(main)/reduce_sum" source_file="<ipython-input-1-5fe284a4623f>" source_line=21}
}

The key thing to notice is that in tst_flat_jit, all array creation happens within a fusion. This essentially means that the compiler is able to deduce the result without actually having to simultaneously allocate space for all of those arrays. In tst_slice_jit, the operation is a bit more involved and the compiler is not able to fuse the entire computation, and so you end up attempting to physically allocate a series of very large arrays on device, the first of which is ~80GB. This leads to a memory error in the second case.

Does that make sense?

[aplund]

the compiler is not able to fuse the entire computation

But the compiler could fuse a computation and infer the result on a slice, right? Is it just that this kind of optimisation isn't something that is done?

[jakevdp]

No, not really: the slice in this case is dynamic, so there's no way for the compiler to know a priori the structure of the particular slice it will be summing over.

[aplund]

The args are explicitly marked as static, and the array is created inside the function. I thought this would make the slice known.

[jakevdp]

Oh, yeah you're right. I suppost the compiler could optimize this, but doesn't currently have the logic for that. It might be interesting to compare the CPU optimization path above to the optimization path on GPU or TPU, which tend to be a bit more performant.