Repository: google/tunix
File: tunix/generate/utils.py (Python)

Code before:
    dst_state: The destination state to transfer to.
    key_mappings: A dictionary defining how to map keys from the source state to
      the target state. The keys of the dictionary are the source keys, and the
      values are tuples containing the target key and the sharding information.
    key_mapping_hook_fns: A dictionary mapping keys to hook functions that
      modify the values before assignment. The hook fn will be called after the
      transpose operation if transpose were to be applied.
    transpose_keys: A dictionary defining which keys to transpose and the
      corresponding axes to transpose.
    reshard_fn: A function to shard the value.

  Returns:
    The target state with the transferred values.
  """
  tgt_flat_list = dst_state.flat_state()
  sharding_dict = None
  if reshard_fn:
    sharding_dict = dict(
        [(key, tgt_params.value.sharding) for key, tgt_params in tgt_flat_list]
    )

  # Maps source keys to target tensor(s) and sharding spec
  src_to_tgt_map = build_flat_dict(tgt_flat_list, key_mappings)

  # Flatten the source state, unrolling any scanned layers.
  unscanned_src_to_tgt_flat = {}
  for src_keys, src_val in src_state.flat_state():
    src_key = '.'.join(str(k) for k in src_keys)
    if src_key not in src_to_tgt_map:
      if 'rng' in src_key:
        logging.debug('Skipping RNG parameter: %s', src_key)
      else:
        logging.error('!!! No mapping for source key: %s', src_key)
      continue

    tgt_param, _, sharding_spec = src_to_tgt_map[src_key]

    def _get_layer_axis_from_sharding_spec(sharding_spec):
      if isinstance(sharding_spec, (list, tuple)):
        for i, spec in enumerate(sharding_spec):
          if spec == 'layer':
            return i
      return None

    layer_axis = _get_layer_axis_from_sharding_spec(sharding_spec)
    if layer_axis is not None:
      num_layers = src_val.value.shape[layer_axis]
      # values = []
      for i in range(num_layers):
        idx = [slice(None)] * src_val.value.ndim
        idx[layer_axis] = i

Review comment (performance):
One performance issue I see with unscanning the layer all together is duplicating the HBM usage. Since JAX slicing creates a new array instead of reference handle. 

Is this something you think we need to worry about in this stage?

Code after:
    dst_state: The destination state to transfer to.
    key_mappings: A dictionary defining how to map keys from the source state to
      the target state. The keys of the dictionary are the source keys, and the
      values are tuples containing the target key and the sharding information.
    key_mapping_hook_fns: A dictionary mapping keys to hook functions that
      modify the values before assignment. The hook fn will be called after the
      transpose operation if transpose were to be applied.
    transpose_keys: A dictionary defining which keys to transpose and the
      corresponding axes to transpose.
    reshard_fn: A function to shard the value.

  Returns:
    The target state with the transferred values.
  """
  tgt_flat_list = dst_state.flat_state()
  sharding_dict = None
  if reshard_fn:
    sharding_dict = dict(
        [(key, tgt_params.value.sharding) for key, tgt_params in tgt_flat_list]
    )

  # Maps source keys to target tensor(s) and sharding spec
  src_to_tgt_map = build_flat_dict(tgt_flat_list, key_mappings)

  # Flatten the source state, unrolling any scanned layers.
  unscanned_src_to_tgt_flat = {}
  for src_keys, src_val in src_state.flat_state():
    src_key = '.'.join(str(k) for k in src_keys)
    if src_key not in src_to_tgt_map:
      if 'rng' in src_key:
        logging.debug('Skipping RNG parameter: %s', src_key)
      else:
        logging.error('!!! No mapping for source key: %s', src_key)
      continue

    tgt_param, tgt_path, sharding_spec = src_to_tgt_map[src_key]

    def _get_layer_axis_from_sharding_spec(sharding_spec):
      if isinstance(sharding_spec, (list, tuple)):
        for i, spec in enumerate(sharding_spec):
          if spec == 'layer':
            return i
      return None

    layer_axis = _get_layer_axis_from_sharding_spec(sharding_spec)
    if layer_axis is not None:
      num_layers = src_val.value.shape[layer_axis]
      for i in range(num_layers):
        idx = [slice(None)] * src_val.value.ndim
        idx[layer_axis] = i
        layer_val = src_val.value[tuple(idx)]