docs: def map(self, func, value_shape=None, dtype=None, with_keys=False):
          """
          Apply an array -> array function across an axis.
  
          Array will be aligned so that the desired set of axes
          are in the keys, which may require a transpose/reshape.
  
          Parameters
          ----------
          func : function
              Function of a single array to apply. If with_keys=True,
              function should be of a (tuple, array) pair.
  
          axis : tuple or int, optional, default=(0,)
              Axis or multiple axes to apply function along.
  
          value_shape : tuple, optional, default=None
              Known shape of values resulting from operation. Only
              valid in spark mode.
  
          dtype : numpy dtype, optional, default=None
              Known shape of dtype resulting from operation. Only
              valid in spark mode.
  
          with_keys : bool, optional, default=False
              Include keys as an argument to the function
          """
          axis = self.baseaxes
  
          if self.mode == 'local':
              axes = sorted(tupleize(axis))
              key_shape = [self.shape[axis] for axis in axes]
              reshaped = self._align(axes, key_shape=key_shape)
  
              if with_keys:
                  keys = zip(*unravel_index(range(prod(key_shape)), key_shape))
                  mapped = asarray(list(map(func, zip(keys, reshaped))))
              else:
                  mapped = asarray(list(map(func, reshaped)))
  
              try:
                  elem_shape = mapped[0].shape
              except:
                  elem_shape = (1,)
  
              expand = list(elem_shape)
              expand = [1] if len(expand) == 0 else expand
  
              # invert the previous reshape operation, using the shape of the map result
              linearized_shape_inv = key_shape + expand
              reordered = mapped.reshape(*linearized_shape_inv)
  
              return self._constructor(reordered, mode=self.mode).__finalize__(self, noprop=('index'))
  
          if self.mode == 'spark':
              expand = lambda x: array(func(x), ndmin=1)
              mapped = self.values.map(expand, axis, value_shape, dtype, with_keys)
              return self._constructor(mapped, mode=self.mode).__finalize__(self, noprop=('index',))