diff --git a/tests/jax_support_test.py b/tests/jax_support_test.py index cc6a66b..ebb1a92 100644 --- a/tests/jax_support_test.py +++ b/tests/jax_support_test.py @@ -12,6 +12,11 @@ # See the License for the specific language governing permissions and # limitations under the License. +import os +import subprocess +import sys +import textwrap + from absl.testing import absltest from absl.testing import parameterized from jax import numpy as jnp @@ -32,6 +37,57 @@ def test_compute_summary(self, dtype): treescope.external.jax_support.summarize_array_data(inp), expected ) + def test_truncate_explicitly_sharded_array(self): + if not hasattr(jax.sharding, "AxisType"): + self.skipTest("JAX does not support explicit sharding.") + + env = os.environ.copy() + env["JAX_PLATFORMS"] = "cpu" + env["XLA_FLAGS"] = "--xla_force_host_platform_device_count=4" + script = textwrap.dedent( + """ + import jax + import numpy as np + import treescope.external.jax_support + + devices = np.asarray(jax.devices()) + mesh = jax.sharding.Mesh( + devices.reshape((1, 4)), + axis_names=("dp", "tp"), + axis_types=( + jax.sharding.AxisType.Explicit, + jax.sharding.AxisType.Explicit, + ), + ) + sharding = jax.sharding.NamedSharding( + mesh, jax.sharding.PartitionSpec("tp", None) + ) + original = np.arange(1600 * 64, dtype=np.int32).reshape((1600, 64)) + array = jax.device_put(original, sharding) + + truncated, valid = ( + treescope.external.jax_support.JAXArrayAdapter() + .get_array_data_with_truncation(array, None, (3, 3)) + ) + + assert truncated.shape == (7, 7) + np.testing.assert_array_equal(truncated[:3, :3], original[:3, :3]) + np.testing.assert_array_equal(truncated[-3:, -3:], original[-3:, -3:]) + expected_valid = np.ones((7, 7), dtype=bool) + expected_valid[3, :] = False + expected_valid[:, 3] = False + np.testing.assert_array_equal(valid, expected_valid) + """ + ) + completed = subprocess.run( + [sys.executable, "-c", script], + check=False, + capture_output=True, + text=True, + env=env, + ) + self.assertEqual(completed.returncode, 0, completed.stderr) + def test_summarize_prng_key(self): keys = jax.random.split(jax.random.key(0, impl="threefry2x32"), 10) summarized = treescope.external.jax_support.summarize_array_data(keys) diff --git a/treescope/external/jax_support.py b/treescope/external/jax_support.py index 6050cfc..17d839c 100644 --- a/treescope/external/jax_support.py +++ b/treescope/external/jax_support.py @@ -117,6 +117,7 @@ def _truncate_part_with_slices( prefix_slices: tuple[slice, ...], remaining_edge_items_per_axis: tuple[int | None, ...], xnp=None, + slice_out_sharding=None, ) -> tuple[jax.Array, jax.Array]: """Helper to truncate names of an array. @@ -128,6 +129,8 @@ def _truncate_part_with_slices( remaining_edge_items_per_axis: Number of edge items to keep for each axis, ignoring any axes whose slices are already computed in `prefix_slices`. xnp: backend to use (numpy or jax.numpy). + slice_out_sharding: Optional replicated sharding to use when converting a + slice on an explicitly-sharded axis into a gather. Returns: Truncated array and mask, which will both be the same shape. @@ -142,7 +145,34 @@ def _truncate_part_with_slices( if not remaining_edge_items_per_axis: # Perform the base case slice. assert len(prefix_slices) == len(array.shape) - truncated_array = array[prefix_slices] + if slice_out_sharding is None: + truncated_array = array[prefix_slices] + else: + first_sliced_axis = next( + ( + axis + for axis, array_slice in enumerate(prefix_slices) + if array_slice != slice(None) + ), + None, + ) + if first_sliced_axis is None: + truncated_array = array + else: + array_slice = prefix_slices[first_sliced_axis] + start, stop, step = array_slice.indices(array.shape[first_sliced_axis]) + indices = xnp.arange(start, stop, step) + gather_index = ( + (slice(None),) * first_sliced_axis + + (indices,) + + (slice(None),) * (array.ndim - first_sliced_axis - 1) + ) + truncated_array = array.at[gather_index].get( + out_sharding=slice_out_sharding + ) + remaining_slices = list(prefix_slices) + remaining_slices[first_sliced_axis] = slice(None) + truncated_array = truncated_array[tuple(remaining_slices)] valid_mask_slices = tuple( slice(None) if mask.shape[i] == 1 else array_slice @@ -165,6 +195,7 @@ def _truncate_part_with_slices( prefix_slices=prefix_slices + (slice(None),), remaining_edge_items_per_axis=remaining_edge_items_per_axis[1:], xnp=xnp, + slice_out_sharding=slice_out_sharding, ) else: assert array.shape[axis] > 2 * edge_items @@ -174,6 +205,7 @@ def _truncate_part_with_slices( prefix_slices=prefix_slices + (slice(None, edge_items),), remaining_edge_items_per_axis=remaining_edge_items_per_axis[1:], xnp=xnp, + slice_out_sharding=slice_out_sharding, ) result_b, valid_b = _truncate_part_with_slices( array, @@ -181,6 +213,7 @@ def _truncate_part_with_slices( prefix_slices=prefix_slices + (slice(-edge_items, None),), remaining_edge_items_per_axis=remaining_edge_items_per_axis[1:], xnp=xnp, + slice_out_sharding=slice_out_sharding, ) padding_shape = list(result_a.shape) padding_shape[axis] = 1 @@ -219,35 +252,50 @@ def truncate_array_and_mask( """ assert jax is not None, "JAX is not available." sharding_kwargs = {} - if hasattr(array, "sharding") and hasattr( - array.sharding, "_device_assignment" - ): + slice_out_sharding = None + if hasattr(array, "sharding"): # _truncate_part_with_slices usually returns slices that have odd - # dimensions, which aren't divisible by most shardings. Unfortunately, - # the XLA GSPMD partitioner sometimes still infers a sharding over one of - # these axes, which then leads to partitioning errors in JAX whenever we - # try to `device_get` the resulting array or call any additional operations - # on it. To avoid this, we'd like to tell JAX to always produce an output - # that is not sharded over any axis. Unfortunately, this is difficult - # because JAX requires the in_shardings and out_shardings to have the same - # devices in the same internal order, and at the time of writing JAX does - # not provide any public API to look up the order of the devices in a - # sharding (it allows looking up the device *set*, but not their order). - # Whether or not this error happens seems to be somewhat nondeterministic. - # To avoid this, we use the private property `_device_assignment` of - # each sharding in order to figure out what device order it has, and then - # explicitly request a fully-replicated output that is definitely safe to - # retrieve. - sharding_kwargs["out_shardings"] = jax.sharding.NamedSharding( - jax.sharding.Mesh(array.sharding._device_assignment, "x"), # pylint: disable=protected-access - jax.sharding.PartitionSpec(), - ) + # dimensions, which aren't divisible by most shardings. Always request a + # fully-replicated output so that the truncated data is safe to retrieve. + if isinstance(array.sharding, jax.sharding.NamedSharding): + output_sharding = jax.sharding.NamedSharding( + array.sharding.mesh, jax.sharding.PartitionSpec() + ) + elif hasattr(array.sharding, "_device_assignment"): + # JAX does not expose the device order for every Sharding implementation, + # so retain the existing private fallback for non-NamedSharding inputs. + output_sharding = jax.sharding.NamedSharding( + jax.sharding.Mesh( + array.sharding._device_assignment, "x" # pylint: disable=protected-access + ), + jax.sharding.PartitionSpec(), + ) + else: + output_sharding = None + + if output_sharding is not None: + sharding_kwargs["out_shardings"] = output_sharding + + # With explicit sharding, a slice itself fails when the sliced dimension + # is not divisible by the mesh axis, before `out_shardings` can take + # effect. Convert the first shrinking slice to a gather with a replicated + # output. + if ( + isinstance(array.sharding, jax.sharding.NamedSharding) + and hasattr(jax.sharding, "AxisType") + and hasattr(array.sharding.mesh, "axis_types") + and jax.sharding.AxisType.Explicit in array.sharding.mesh.axis_types + ): + slice_out_sharding = output_sharding + if array.size < SUMMARIZE_USING_NUMPY_THRESHOLD and safe_to_summarize(array): fn = functools.partial(_truncate_part_with_slices, xnp=np) else: - fn = jax.jit( - _truncate_part_with_slices, static_argnums=(2, 3), **sharding_kwargs + truncate_part = functools.partial( + _truncate_part_with_slices, + slice_out_sharding=slice_out_sharding, ) + fn = jax.jit(truncate_part, static_argnums=(2, 3), **sharding_kwargs) return fn(array, mask, (), edge_items_per_axis)